add

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>

.buildkite/nightly-release-pipeline.yaml

@@ -0,0 +1,20 @@
+steps:
+  - label: "Build and upload wheel"
+    key: "build-wheel"
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - "curl -LsSf https://astral.sh/uv/install.sh | sh"
+      - 'export PATH="$HOME/.local/bin:$PATH"'
+      - "uv venv --python=3.12 && source .venv/bin/activate"
+      - "uv pip install --upgrade build"
+      - "python3 -m build"
+      - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+
+  - label: "Generate and upload wheel indices"
+    depends_on: "build-wheel"
+    allow_dependency_failure: true
+    agents:
+      queue: small_cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"

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

@@ -0,0 +1,87 @@
+#!/usr/bin/env bash
+
+set -ex
+
+# Generate and upload wheel indices for all vllm-omni wheels in the commit directory.
+# This script should run once after all wheels have been built and uploaded.
+# All paths are under the omni/ prefix in the vllm-wheels S3 bucket.
+
+# ======== setup ========
+
+BUCKET="vllm-wheels"
+INDICES_OUTPUT_DIR="indices"
+PYTHON="${PYTHON_PROG:-python3}"
+SUBPATH="omni/$BUILDKITE_COMMIT"
+S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
+
+# detect if python3.12+ is available
+has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
+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"
+fi
+
+echo "Using python interpreter: $PYTHON"
+echo "Python version: $($PYTHON --version)"
+
+# ======== generate and upload indices ========
+
+# list all wheels in the commit directory
+echo "Existing wheels on S3:"
+aws s3 ls "$S3_COMMIT_PREFIX" || echo "(no objects found)"
+obj_json="objects.json"
+aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
+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
+
+# Generate indices -- the version is just the commit hash (not omni/{commit})
+# because relative paths are computed between the index and wheel directories,
+# both of which live under the omni/ prefix in S3.
+$PYTHON .buildkite/scripts/generate-nightly-index.py \
+    --version "$BUILDKITE_COMMIT" \
+    --current-objects "$obj_json" \
+    --output-dir "$INDICES_OUTPUT_DIR" \
+    --comment "commit $BUILDKITE_COMMIT"
+
+# copy indices to /omni/{commit}/ unconditionally
+echo "Uploading indices to $S3_COMMIT_PREFIX"
+aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
+
+# copy to /omni/nightly/ when NIGHTLY=1
+if [[ "${NIGHTLY:-}" == "1" ]]; then
+    echo "Uploading indices to overwrite /omni/nightly/"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/nightly/"
+fi
+
+# detect version from any wheel in the commit directory
+first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].endswith('.whl')), ''))")
+if [[ -z "$first_wheel_key" ]]; then
+    echo "Error: No wheels found in $S3_COMMIT_PREFIX"
+    exit 1
+fi
+first_wheel=$(basename "$first_wheel_key")
+aws s3 cp "s3://$BUCKET/${first_wheel_key}" "/tmp/${first_wheel}"
+version=$(unzip -p "/tmp/${first_wheel}" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
+rm -f "/tmp/${first_wheel}"
+echo "Version in wheel: $version"
+pure_version="${version%%+*}"
+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/"
+    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" \
+        --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/"
+fi

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

@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import json
+import sys
+from dataclasses import asdict, dataclass
+from datetime import datetime
+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."""
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+if not sys.version_info >= (3, 12):
+    raise RuntimeError("This script requires Python 3.12 or higher.")
+
+INDEX_HTML_TEMPLATE = """<!DOCTYPE html>
+<html>
+  <!-- {comment} -->
+  <meta name="pypi:repository-version" content="1.0">
+  <body>
+{items}
+  </body>
+</html>
+"""
+
+
+@dataclass
+class WheelFileInfo:
+    package_name: str
+    version: str
+    build_tag: str | None
+    python_tag: str
+    abi_tag: str
+    platform_tag: str
+    filename: str
+
+
+def parse_from_filename(file: str) -> WheelFileInfo:
+    """
+    Parse wheel filename per PEP 427:
+        {package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl
+    """
+    wheel_file_re = re.compile(
+        r"^(?P<package_name>.+)-(?P<version>[^-]+?)(-(?P<build_tag>[^-]+))?-(?P<python_tag>[^-]+)-(?P<abi_tag>[^-]+)-(?P<platform_tag>[^-]+)\.whl$"
+    )
+    match = wheel_file_re.match(file)
+    if not match:
+        raise ValueError(f"Invalid wheel file name: {file}")
+
+    return WheelFileInfo(
+        package_name=match.group("package_name"),
+        version=match.group("version"),
+        build_tag=match.group("build_tag"),
+        python_tag=match.group("python_tag"),
+        abi_tag=match.group("abi_tag"),
+        platform_tag=match.group("platform_tag"),
+        filename=file,
+    )
+
+
+def generate_project_list(package_names: list[str], comment: str = "") -> str:
+    """Generate top-level PEP 503 project list HTML."""
+    href_tags = []
+    for name in sorted(package_names):
+        href_tags.append(f'    <a href="{name}/">{name}/</a><br/>')
+    return INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
+
+
+def generate_package_index(
+    wheel_files: list[WheelFileInfo],
+    wheel_base_dir: Path,
+    index_base_dir: Path,
+    comment: str = "",
+) -> tuple[str, str]:
+    """Generate package index HTML and metadata JSON linking to wheel files."""
+    href_tags = []
+    metadata = []
+    for file in sorted(wheel_files, key=lambda x: x.filename):
+        relative_path = wheel_base_dir.relative_to(index_base_dir, walk_up=True) / file.filename
+        # handle '+' in URL; avoid double-encoding '/' and '%2B' (AWS S3 behavior)
+        file_path_quoted = quote(relative_path.as_posix(), safe=":%/")
+        href_tags.append(f'    <a href="{file_path_quoted}">{file.filename}</a><br/>')
+        file_meta = asdict(file)
+        file_meta["path"] = file_path_quoted
+        metadata.append(file_meta)
+    index_str = INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
+    metadata_str = json.dumps(metadata, indent=2)
+    return index_str, metadata_str
+
+
+def generate_index(
+    whl_files: list[str],
+    wheel_base_dir: Path,
+    index_base_dir: Path,
+    comment: str = "",
+):
+    """
+    Generate PEP 503 index for all wheel files.
+
+    Output structure:
+        index_base_dir/
+            index.html          # project list linking to vllm-omni/
+            vllm-omni/
+                index.html      # package index linking to wheel files
+                metadata.json   # machine-readable metadata
+    """
+    parsed_files = [parse_from_filename(f) for f in whl_files]
+
+    if not parsed_files:
+        print("No wheel files found, skipping index generation.")
+        return
+
+    comment_str = f" ({comment})" if comment else ""
+    comment_tmpl = f"Generated on {datetime.now().isoformat()}{comment_str}"
+
+    # Group by normalized package name
+    packages: dict[str, list[WheelFileInfo]] = {}
+    for file in parsed_files:
+        name = normalize_package_name(file.package_name)
+        packages.setdefault(name, []).append(file)
+
+    print(f"Found packages: {list(packages.keys())}")
+
+    # Generate per-package index
+    for package, files in packages.items():
+        package_dir = index_base_dir / package
+        package_dir.mkdir(parents=True, exist_ok=True)
+        index_str, metadata_str = generate_package_index(files, wheel_base_dir, package_dir, comment)
+        with open(package_dir / "index.html", "w") as f:
+            f.write(index_str)
+        with open(package_dir / "metadata.json", "w") as f:
+            f.write(metadata_str)
+
+    # Generate top-level project list
+    project_list_str = generate_project_list(sorted(packages.keys()), comment_tmpl)
+    with open(index_base_dir / "index.html", "w") as f:
+        f.write(project_list_str)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Generate PEP 503 wheel index from S3 object listing.")
+    parser.add_argument("--version", type=str, required=True, help="Version string (e.g., commit hash)")
+    parser.add_argument("--current-objects", type=str, required=True, help="Path to JSON from S3 list-objects-v2")
+    parser.add_argument("--output-dir", type=str, required=True, help="Directory to write index files")
+    parser.add_argument("--wheel-dir", type=str, default=None, help="Wheel directory (defaults to --version)")
+    parser.add_argument("--comment", type=str, default="", help="Comment for generated HTML")
+
+    args = parser.parse_args()
+
+    version = args.version
+    if "\\" in version or "/" in version:
+        raise ValueError("Version string must not contain slashes or backslashes.")
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    with open(args.current_objects) as f:
+        current_objects: dict[str, list[dict[str, Any]]] = json.load(f)
+
+    wheel_files = [
+        item["Key"].split("/")[-1] for item in current_objects.get("Contents", []) if item["Key"].endswith(".whl")
+    ]
+
+    print(f"Found {len(wheel_files)} wheel files for version {version}: {wheel_files}")
+
+    # For release versions, filter to only matching non-dev wheels
+    PY_VERSION_RE = re.compile(r"^\d+\.\d+\.\d+([a-zA-Z0-9.+-]*)?$")
+    if PY_VERSION_RE.match(version):
+        wheel_files = [f for f in wheel_files if version in f and "dev" not in f]
+        print(f"Non-nightly version detected, wheel files used: {wheel_files}")
+    else:
+        print("Nightly version detected, keeping all wheel files.")
+
+    wheel_dir = (args.wheel_dir or version).strip().rstrip("/")
+    wheel_base_dir = Path(output_dir).parent / wheel_dir
+    index_base_dir = Path(output_dir)
+
+    generate_index(
+        whl_files=wheel_files,
+        wheel_base_dir=wheel_base_dir,
+        index_base_dir=index_base_dir,
+        comment=args.comment.strip(),
+    )
+    print(f"Successfully generated index in {output_dir}")

.buildkite/scripts/upload-nightly-wheels.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+
+set -ex
+
+# Upload a single wheel to S3 under the omni/ prefix.
+# Index generation is handled separately by generate-and-upload-nightly-index.sh.
+
+BUCKET="vllm-wheels"
+SUBPATH="omni/$BUILDKITE_COMMIT"
+S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
+
+# ========= collect & upload the wheel ==========
+
+# python3 -m build outputs to dist/ by default
+wheel_files=(dist/*.whl)
+
+# Check that exactly one wheel is found
+if [[ ${#wheel_files[@]} -ne 1 ]]; then
+  echo "Error: Expected exactly one wheel file in dist/, but found ${#wheel_files[@]}"
+  exit 1
+fi
+wheel="${wheel_files[0]}"
+
+echo "Uploading wheel: $wheel"
+
+# Extract the version from the wheel
+version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
+echo "Version in wheel: $version"
+
+# Upload wheel to S3
+aws s3 cp "$wheel" "$S3_COMMIT_PREFIX"
+
+echo "Wheel uploaded to $S3_COMMIT_PREFIX. Index generation is handled by a separate step."

.buildkite/test-merge.yml

@@ -259,7 +259,6 @@ steps:
     depends_on: upload-merge-pipeline
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
       - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
       - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
       - pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"

.buildkite/test-nightly.yml

@@ -355,6 +355,48 @@ steps:
                   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:c392ce21e9cf9ea65c52b866447793db10e0261c
+                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: Qwen-Image Diffusion Perf Test with H100"
     key: nightly-qwen-image-performance
     timeout_in_minutes: 300

.buildkite/test-ready.yml

@@ -216,44 +216,42 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  # - label: "Omni Model Test with H100"
-  #   depends_on: upload-ready-pipeline
-  #   commands:
-  #     - |
-  #       timeout 20m bash -c '
-  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #         export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-  #         # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-  #         pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_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
-  #           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: "Omni Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_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
+            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: "MiMo-Audio E2E Test with H100"
     depends_on: upload-ready-pipeline

README.md

@@ -16,7 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
-
+- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
 - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
 - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
 - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.

docs/community/governance.md

@@ -37,10 +37,12 @@ Committers have write access and merge rights. They typically have deep expertis
 - [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
 - [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
 - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
+- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
 - [@princepride](https://github.com/princepride): Diffusion and Omni Support
 - [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
 - [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
-- [@wtomin](https://github.com/wtomin):
+- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
+- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
 - [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
 
 ## Meetings

docs/design/module/dit_module.md

@@ -192,7 +192,7 @@ class _BaseScheduler(SchedulerInterface):
         self._waiting = deque()
         self._running = []
         self._finished_req_ids = set()
-        self._max_batch_size = 1
+        self.max_num_running_reqs = 1
 ```
 
 **Design Features**:
@@ -201,7 +201,7 @@ class _BaseScheduler(SchedulerInterface):
 
 - **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
 
-- **Current constraint**: `_max_batch_size` remains `1` because the current engine path is still synchronous request-mode execution.
+- **Current constraint**: `max_num_running_reqs` remains `1` because the current engine path is still synchronous request-mode execution.
 
 #### 2.4 Current `RequestScheduler` Policy
 

docs/user_guide/diffusion_features.md

@@ -15,6 +15,7 @@ vLLM-Omni supports various advanced features for diffusion models:
 - Acceleration: **cache methods**, **parallelism methods**
 - Memory optimization: **cpu offloading**, **quantization**
 - Extensions: **LoRA inference**
+- Execution modes: **step execution**
 
 ## Supported Features
 
@@ -64,6 +65,16 @@ Extension methods add specialized capabilities to diffusion models beyond standa
 | **[LoRA Inference](diffusion/lora.md)** | Enables inference with Low-Rank Adaptation (LoRA) adapters weights | Reinforcement learning extensions |
 
 
+### Execution Modes
+
+Execution modes control how the diffusion pipeline processes denoise steps.
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[Step Execution](diffusion/step_execution.md)** | Per-step denoise execution with mid-request abort support | Request cancellation between denoise steps, fine-grained execution control |
+
+**Note:** Step execution is currently supported by QwenImagePipeline only. See [Supported Models](#supported-models) for details.
+
 ### Quantization Methods
 
 | Method | Configuration | Description | Best For |
@@ -87,28 +98,28 @@ The following tables show which models support each feature:
 
 ### ImageGen
 
-| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
-|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
-| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
-| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
-| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
-| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ |
-| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
-| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ❌ | ✅ | ✅ |
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
+| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
+| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
+| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
+| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ❌ | ✅ | ✅ | ❌ |
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
@@ -116,19 +127,19 @@ The following tables show which models support each feature:
 
 ### VideoGen
 
-| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
-|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
-| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 
 ### AudioGen
 
-| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
-|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
-| **Stable-Audio-Open** | ❌ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ |
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Stable-Audio-Open** | ❌ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
 
 
 ## Feature Compatibility
@@ -139,21 +150,22 @@ The following tables show which models support each feature:
 - ❌: No support plan
 - ❓: Not verified yet and Not Recommended
 
-|  | ⚡TeaCache | ⚡Cache-DiT | 🔀Ulysses-SP | 🔀Ring-Attn | 🔀CFG-Parallel | 🔀Tensor Parallel | 🔀HSDP | 🔀Expert Parallel | 💾CPU Offloading (Layerwise) | 💾CPU Offloading (Module-wise) | 💾VAE Patch Parallel | 💾FP8 Quant | 🔧LoRA Inference |
-|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
-| **⚡TeaCache** | | | | | | | | | | | | | |
-| **⚡Cache-DiT** | ❌ | | | | | | | | | | | | |
-| **🔀Ulysses-SP** | ✅ | ✅ | | | | | | | | | | | |
-| **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | |
-| **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
-| **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | |
-| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | |
-| **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | |
-| **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | |
-| **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | |
-| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | |
-| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | |
-| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | |
+|  | ⚡TeaCache | ⚡Cache-DiT | 🔀Ulysses-SP | 🔀Ring-Attn | 🔀CFG-Parallel | 🔀Tensor Parallel | 🔀HSDP | 🔀Expert Parallel | 💾CPU Offloading (Layerwise) | 💾CPU Offloading (Module-wise) | 💾VAE Patch Parallel | 💾FP8 Quant | 🔧LoRA Inference | 🔄Step Execution |
+|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| **⚡TeaCache** | | | | | | | | | | | | | | |
+| **⚡Cache-DiT** | ❌ | | | | | | | | | | | | | |
+| **🔀Ulysses-SP** | ✅ | ✅ | | | | | | | | | | | | |
+| **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | | |
+| **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
+| **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
+| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | | |
+| **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | | |
+| **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | | |
+| **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | | |
+| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | | |
+| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | | |
+| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | |
+| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ❌ | |
 
 !!! info
 
@@ -162,6 +174,7 @@ The following tables show which models support each feature:
     3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
     4. CPU Offloading (Layerwise) supports single-card for now.
     5. Using FP8-Quant as an example of qunatization methods.
+    6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
 
 
 ## Learn More
@@ -185,6 +198,10 @@ The following tables show which models support each feature:
 
 - **[LoRA Inference Guide](diffusion/lora.md)** - Low-Rank Adaptation for style customization and fine-tuning
 
+**Execution Modes:**
+
+- **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support
+
 **Advanced Topics:**
 
 - **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance

docs/user_guide/examples/online_serving/image_to_video.md

@@ -37,12 +37,45 @@ artifact, poll the job status and then download the completed file from the
 content endpoint.
 
 The main endpoints are:
-- `POST /v1/videos`: create a video generation job
+- `POST /v1/videos`: create a video generation job (async)
+- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
 - `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
 - `GET /v1/videos`: list stored video jobs
 - `GET /v1/videos/{video_id}/content`: download the generated video file
 - `DELETE /v1/videos/{video_id}`: delete the job and any stored output
 
+## Sync API (Benchmark / Testing)
+
+`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
+completes and returns the raw video bytes (`video/mp4`) directly in the response
+body. It is designed for benchmark and testing scenarios where one-shot
+request/response latency measurement is needed.
+
+The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
+not create any stored job record — the response is purely the generated video
+file. Metadata is returned via response headers:
+
+- `X-Request-Id`: unique identifier for this generation request
+- `X-Model`: model name used for generation
+- `X-Inference-Time-S`: wall-clock inference time in seconds
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "input_reference=@/path/to/input.png" \
+  -F "size=832x480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=12.0" \
+  -F "seed=42" \
+  -o sync_i2v_output.mp4
+```
+
 ## Storage
 
 Generated video files are stored on local disk by the async video API.

docs/user_guide/examples/online_serving/text_to_video.md

@@ -37,12 +37,43 @@ artifact, poll the job status and then download the completed file from the
 content endpoint.
 
 The main endpoints are:
-- `POST /v1/videos`: create a video generation job
+- `POST /v1/videos`: create a video generation job (async)
+- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
 - `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
 - `GET /v1/videos`: list stored video jobs
 - `GET /v1/videos/{video_id}/content`: download the generated video file
 - `DELETE /v1/videos/{video_id}`: delete the job and any stored output
 
+## Sync API (Benchmark / Testing)
+
+`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
+completes and returns the raw video bytes (`video/mp4`) directly in the response
+body. It is designed for benchmark and testing scenarios where one-shot
+request/response latency measurement is needed.
+
+The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
+not create any stored job record — the response is purely the generated video
+file. Metadata is returned via response headers:
+
+- `X-Request-Id`: unique identifier for this generation request
+- `X-Model`: model name used for generation
+- `X-Inference-Time-S`: wall-clock inference time in seconds
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+  -F "size=832x480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.0" \
+  -F "guidance_scale_2=4.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=5.0" \
+  -F "seed=42" \
+  -o sync_t2v_output.mp4
+```
+
 ## Storage
 
 Generated video files are stored on local disk by the async video API.

examples/offline_inference/bagel/end2end.py

@@ -101,7 +101,7 @@ def main():
 
     if not prompts:
         # Default prompt for text2img test if none provided
-        prompts = ["<|im_start|>A cute cat<|im_end|>"]
+        prompts = ["A cute cat"]
         print(f"[Info] No prompts provided, using default: {prompts}")
     omni_outputs = []
 
@@ -168,7 +168,6 @@ def main():
 
     params_list = omni.default_sampling_params_list
     if args.modality in ("text2img", "img2img"):
-        params_list[0].max_tokens = 1  # type: ignore
         if len(params_list) > 1:
             diffusion_params = params_list[1]
             diffusion_params.num_inference_steps = args.steps  # type: ignore

examples/online_serving/bagel/openai_chat_client.py

@@ -125,7 +125,7 @@ def generate_image(
 
 def main():
     parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
-    parser.add_argument("--prompt", "-p", default="<|im_start|>A cute cat<|im_end|>", help="Text prompt")
+    parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
     parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
     parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
 

examples/online_serving/image_to_video/README.md

@@ -34,12 +34,45 @@ artifact, poll the job status and then download the completed file from the
 content endpoint.
 
 The main endpoints are:
-- `POST /v1/videos`: create a video generation job
+- `POST /v1/videos`: create a video generation job (async)
+- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
 - `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
 - `GET /v1/videos`: list stored video jobs
 - `GET /v1/videos/{video_id}/content`: download the generated video file
 - `DELETE /v1/videos/{video_id}`: delete the job and any stored output
 
+## Sync API (Benchmark / Testing)
+
+`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
+completes and returns the raw video bytes (`video/mp4`) directly in the response
+body. It is designed for benchmark and testing scenarios where one-shot
+request/response latency measurement is needed.
+
+The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
+not create any stored job record — the response is purely the generated video
+file. Metadata is returned via response headers:
+
+- `X-Request-Id`: unique identifier for this generation request
+- `X-Model`: model name used for generation
+- `X-Inference-Time-S`: wall-clock inference time in seconds
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "input_reference=@/path/to/input.png" \
+  -F "size=832x480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=12.0" \
+  -F "seed=42" \
+  -o sync_i2v_output.mp4
+```
+
 ## Storage
 
 Generated video files are stored on local disk by the async video API.

examples/online_serving/qwen3_tts/gradio_demo.py

@@ -125,7 +125,7 @@ PLAYER_HTML = """
 def _build_player_js(sample_rate: int) -> str:
     """Build the JavaScript that powers the AudioWorklet player."""
     return f"""
-() => {{
+    <script>
     const SR = {sample_rate};
     const WC = {json.dumps(WORKLET_JS)};
     let ctx = null, node = null, abort = null, gen = false, st = {{}};
@@ -284,7 +284,7 @@ def _build_player_js(sample_rate: int) -> str:
             }}
         }}
     }};
-}}
+    </script>
 """
 
 
@@ -435,10 +435,7 @@ def create_app(api_base: str):
     )
 
     with gr.Blocks(
-        css=css,
         title="Qwen3-TTS Demo",
-        js=_build_player_js(PCM_SAMPLE_RATE),
-        theme=theme,
     ) as demo:
         gr.HTML(f"""
         <div style="display:flex; align-items:center; gap:16px; margin-bottom:8px;">
@@ -773,7 +770,14 @@ def create_app(api_base: str):
 
         demo.queue()
 
-    return gr.mount_gradio_app(fastapi_app, demo, path="/")
+    return gr.mount_gradio_app(
+        fastapi_app,
+        demo,
+        path="/",
+        css=css,
+        theme=theme,
+        head=_build_player_js(PCM_SAMPLE_RATE),
+    )
 
 
 def parse_args():

examples/online_serving/qwen3_tts/tts_common.py

@@ -103,7 +103,7 @@ def build_payload(
 
     if task_type == "CustomVoice":
         if voice:
-            payload["speaker"] = voice
+            payload["voice"] = voice
         if instructions and instructions.strip():
             payload["instructions"] = instructions.strip()
 

examples/online_serving/text_to_video/README.md

@@ -34,12 +34,43 @@ artifact, poll the job status and then download the completed file from the
 content endpoint.
 
 The main endpoints are:
-- `POST /v1/videos`: create a video generation job
+- `POST /v1/videos`: create a video generation job (async)
+- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
 - `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
 - `GET /v1/videos`: list stored video jobs
 - `GET /v1/videos/{video_id}/content`: download the generated video file
 - `DELETE /v1/videos/{video_id}`: delete the job and any stored output
 
+## Sync API (Benchmark / Testing)
+
+`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
+completes and returns the raw video bytes (`video/mp4`) directly in the response
+body. It is designed for benchmark and testing scenarios where one-shot
+request/response latency measurement is needed.
+
+The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
+not create any stored job record — the response is purely the generated video
+file. Metadata is returned via response headers:
+
+- `X-Request-Id`: unique identifier for this generation request
+- `X-Model`: model name used for generation
+- `X-Inference-Time-S`: wall-clock inference time in seconds
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+  -F "size=832x480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.0" \
+  -F "guidance_scale_2=4.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=5.0" \
+  -F "seed=42" \
+  -o sync_t2v_output.mp4
+```
+
 ## Storage
 
 Generated video files are stored on local disk by the async video API.

pyproject.toml

@@ -177,20 +177,24 @@ markers = [
     "rocm: Tests that run on AMD/ROCm (auto-added)",
     "xpu: Tests that run on XPU (auto-added)",
     "npu: Tests that run on NPU/Ascend (auto-added)",
+    "musa: Tests that run on MUSA/Moore Threads (auto-added)",
     # specified computation resources marks (auto-added)
     "H100: Tests that require H100 GPU",
     "L4: Tests that require L4 GPU",
     "MI325: Tests that require MI325 GPU (AMD/ROCm)",
+    "S5000: Tests that require S5000 GPU (Moore Threads/MUSA)",
     "A2: Tests that require A2 NPU",
     "A3: Tests that require A3 NPU",
     "distributed_cuda: Tests that require multi cards on CUDA platform",
     "distributed_rocm: Tests that require multi cards on ROCm platform",
     "distributed_xpu: Tests that require multi cards on XPU platform",
     "distributed_npu: Tests that require multi cards on NPU platform",
+    "distributed_musa: Tests that require multi cards on MUSA platform",
     "skipif_cuda: Skip if the num of CUDA cards is less than the required",
     "skipif_rocm: Skip if the num of ROCm cards is less than the required",
     "skipif_xpu: Skip if the num of XPU cards is less than the required",
     "skipif_npu: Skip if the num of NPU cards is less than the required",
+    "skipif_musa: Skip if the num of MUSA cards is less than the required",
     # more detailed markers
     "slow: Slow tests (may skip in quick CI)",
     "benchmark: Benchmark tests",

requirements/musa.txt

@@ -0,0 +1,4 @@
+-r common.txt
+# MUSA platform dependencies
+torchada>=0.1.46
+onnxruntime>=1.23.2

setup.py

@@ -3,7 +3,7 @@ Setup script for vLLM-Omni with hardware-dependent installation.
 
 This setup.py implements platform-aware dependency routing so users can run
 `pip install vllm-omni` and automatically receive the correct platform-specific
-dependencies (CUDA/ROCm/CPU/XPU/NPU) without requiring extras like `[cuda]`.
+dependencies (CUDA/ROCm/CPU/XPU/NPU/MUSA) without requiring extras like `[cuda]`.
 """
 
 import os
@@ -46,16 +46,16 @@ def detect_target_device() -> str:
 
     Priority order:
     1. VLLM_OMNI_TARGET_DEVICE environment variable (highest priority)
-    2. Torch backend detection (cuda, rocm, npu, xpu)
+    2. Torch backend detection (cuda, rocm, npu, xpu, musa)
     3. CPU fallback (default)
 
     Returns:
-        str: Device name ('cuda', 'rocm', 'npu', 'xpu', or 'cpu')
+        str: Device name ('cuda', 'rocm', 'npu', 'xpu', 'musa', or 'cpu')
     """
     # Priority 1: Explicit override via environment variable
     target_device = os.environ.get("VLLM_OMNI_TARGET_DEVICE")
     if target_device:
-        valid_devices = ["cuda", "rocm", "npu", "xpu", "cpu"]
+        valid_devices = ["cuda", "rocm", "npu", "xpu", "musa", "cpu"]
         if target_device.lower() in valid_devices:
             print(f"Using target device from VLLM_OMNI_TARGET_DEVICE: {target_device.lower()}")
             return target_device.lower()
@@ -97,6 +97,15 @@ def detect_target_device() -> str:
             except Exception:
                 pass
 
+        # Check for MUSA (Moore Threads)
+        if hasattr(torch, "musa"):
+            try:
+                if torch.musa.is_available():
+                    print("Detected MUSA backend from torch")
+                    return "musa"
+            except Exception:
+                pass
+
         print("No GPU backend detected in torch, defaulting to CPU")
         return "cpu"
 
@@ -152,6 +161,8 @@ def get_vllm_omni_version() -> str:
             version += f"{sep}npu"
         elif device == "xpu":
             version += f"{sep}xpu"
+        elif device == "musa":
+            version += f"{sep}musa"
         elif device == "cpu":
             version += f"{sep}cpu"
         else:

tests/conftest.py

@@ -48,6 +48,7 @@ from vllm.utils.network_utils import get_open_port
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -688,8 +689,112 @@ def generate_synthetic_audio(
     return result
 
 
-def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_file: bool = False) -> dict[str, Any]:
-    """Generate synthetic video with bouncing balls and return base64 string."""
+def _mux_mp4_bytes_with_synthetic_audio(
+    video_mp4_bytes: bytes,
+    *,
+    num_frames: int,
+    fps: float = 30.0,
+    sample_rate: int = 48000,
+) -> bytes:
+    """
+    Mux a video-only MP4 with mono TTS audio from :func:`generate_synthetic_audio` (AAC).
+
+    Audio length is at least the video duration in whole seconds (rounded up); ffmpeg
+    ``-shortest`` trims to the video when the WAV is longer.
+
+    Uses ffmpeg from ``imageio_ffmpeg`` when available, else ``ffmpeg`` on PATH.
+    If TTS or mux fails, returns ``video_mp4_bytes`` unchanged.
+
+    Mux subprocess does **not** use ``capture_output=True``: ffmpeg can block writing
+    to a full stderr pipe while :func:`subprocess.run` waits for exit (classic deadlock).
+    """
+    duration_sec = num_frames / fps if fps > 0 else 0.0
+    # generate_synthetic_audio(duration=int) uses at least 1s of buffer internally
+    duration_int = max(1, int(math.ceil(duration_sec)))
+
+    try:
+        audio_result = generate_synthetic_audio(
+            duration=duration_int,
+            num_channels=1,
+            sample_rate=sample_rate,
+            save_to_file=False,
+        )
+        audio_pcm = audio_result["np_array"]
+    except Exception as e:
+        logger.warning("Synthetic video: generate_synthetic_audio failed (%s); using video-only MP4.", e)
+        return video_mp4_bytes
+
+    try:
+        import imageio_ffmpeg
+
+        ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
+    except Exception:
+        ffmpeg_exe = "ffmpeg"
+
+    import tempfile
+
+    try:
+        with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp:
+            vid_path = os.path.join(tmp, "video.mp4")
+            wav_path = os.path.join(tmp, "audio.wav")
+            out_path = os.path.join(tmp, "out.mp4")
+            with open(vid_path, "wb") as f:
+                f.write(video_mp4_bytes)
+            sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16")
+            cmd = [
+                ffmpeg_exe,
+                "-y",
+                "-nostdin",
+                "-hide_banner",
+                "-loglevel",
+                "error",
+                "-i",
+                vid_path,
+                "-i",
+                wav_path,
+                "-c:v",
+                "copy",
+                "-c:a",
+                "aac",
+                "-b:a",
+                "128k",
+                "-shortest",
+                "-movflags",
+                "+faststart",
+                out_path,
+            ]
+            subprocess.run(
+                cmd,
+                check=True,
+                stdin=subprocess.DEVNULL,
+                timeout=300,
+            )
+            with open(out_path, "rb") as f:
+                return f.read()
+    except (
+        FileNotFoundError,
+        subprocess.CalledProcessError,
+        subprocess.TimeoutExpired,
+        OSError,
+    ) as e:
+        logger.warning("Synthetic video: audio mux failed (%s); using video-only MP4.", e)
+        return video_mp4_bytes
+
+
+def generate_synthetic_video(
+    width: int,
+    height: int,
+    num_frames: int,
+    save_to_file: bool = False,
+    *,
+    embed_audio: bool = False,
+) -> dict[str, Any]:
+    """Generate synthetic video with bouncing balls and base64 MP4.
+
+    When ``embed_audio`` is True, muxes mono AAC from :func:`generate_synthetic_audio`
+    (TTS + ffmpeg) into the MP4; otherwise returns video-only MP4 (faster when tests do
+    not need an audio track).
+    """
 
     import cv2
     import imageio
@@ -762,13 +867,13 @@ def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_f
     result = {
         "np_array": video_array,
     }
-    video_bytes = None
     saved_file_path = None
 
+    fps = 30
     buffer = io.BytesIO()
     writer_kwargs = {
         "format": "mp4",
-        "fps": 30,
+        "fps": fps,
         "codec": "libx264",
         "quality": 7,
         "pixelformat": "yuv420p",
@@ -787,32 +892,31 @@ def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_f
         ],
     }
 
-    if save_to_file:
-        import datetime
-
-        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        output_path = f"video_{width}x{height}_{timestamp}.mp4"
-        try:
-            with imageio.get_writer(output_path, **writer_kwargs) as writer:
-                for frame in video_frames:
-                    writer.append_data(frame)
-
-            saved_file_path = output_path
-            print(f"Video saved to: {saved_file_path}")
-            with open(output_path, "rb") as f:
-                video_bytes = f.read()
-
-        except Exception as e:
-            print(f"Warning: Failed to save video to file {output_path}: {e}")
-            save_to_file = False
-
-    if not save_to_file or video_bytes is None:
+    try:
         with imageio.get_writer(buffer, **writer_kwargs) as writer:
             for frame in video_frames:
                 writer.append_data(frame)
-
         buffer.seek(0)
-        video_bytes = buffer.read()
+        video_only_bytes = buffer.read()
+    except Exception as e:
+        print(f"Warning: Failed to encode synthetic video: {e}")
+        raise
+
+    if embed_audio:
+        video_bytes = _mux_mp4_bytes_with_synthetic_audio(video_only_bytes, num_frames=num_frames, fps=float(fps))
+    else:
+        video_bytes = video_only_bytes
+
+    if save_to_file:
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        output_path = f"video_{width}x{height}_{timestamp}.mp4"
+        try:
+            with open(output_path, "wb") as f:
+                f.write(video_bytes)
+            saved_file_path = output_path
+            print(f"Video saved to: {saved_file_path}")
+        except Exception as e:
+            print(f"Warning: Failed to save video to file {output_path}: {e}")
 
     base64_video = base64.b64encode(video_bytes).decode("utf-8")
 
@@ -962,7 +1066,7 @@ def convert_audio_to_text(audio_data):
     Convert base64 encoded audio data to text using speech recognition.
     """
     audio_data = base64.b64decode(audio_data)
-    output_path = f"./test_{int(time.time())}.wav"
+    output_path = f"./test_{uuid.uuid4().hex}.wav"
     with open(output_path, "wb") as audio_file:
         audio_file.write(audio_data)
 
@@ -986,8 +1090,24 @@ def _merge_base64_audio_to_segment(base64_list: list[str]):
 def _whisper_transcribe_in_current_process(output_path: str) -> str:
     import whisper
 
-    # Keep Whisper on CPU to avoid consuming GPU memory in tests.
-    model = whisper.load_model("small", device="cpu")
+    # Multi-GPU: use last visible device to avoid colliding with default device 0; single device uses 0.
+    device_index = None
+    if current_omni_platform.is_available():
+        n = current_omni_platform.get_device_count()
+        if n == 1:
+            device_index = 0
+        elif n > 1:
+            device_index = n - 1
+
+    if device_index is not None:
+        torch_device = current_omni_platform.get_torch_device(device_index)
+        current_omni_platform.set_device(torch_device)
+        device = str(torch_device)
+        use_accelerator = True
+    else:
+        use_accelerator = False
+        device = "cpu"
+    model = whisper.load_model("small", device=device)
     try:
         text = model.transcribe(
             output_path,
@@ -998,6 +1118,9 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
     finally:
         del model
         gc.collect()
+        if use_accelerator:
+            current_omni_platform.synchronize()
+            current_omni_platform.empty_cache()
 
     return text or ""
 
@@ -1262,10 +1385,12 @@ def modify_stage_config(
                 # Direct top-level key
                 config[key] = value
 
-    # Save to new file with timestamp
-    timestamp = int(time.time())
+    # Unique suffix: multiple modify_stage_config calls in one process often run
+    # within the same second (e.g. test_qwen3_omni_expansion imports both
+    # get_chunk_config and get_batch_token_config). int(time.time()) would collide
+    # and the later write would overwrite the earlier YAML on disk.
     base_name = yaml_path.rsplit(".", 1)[0] if "." in yaml_path else yaml_path
-    output_path = f"{base_name}_{timestamp}.yaml"
+    output_path = f"{base_name}_{time.time_ns()}.yaml"
 
     with open(output_path, "w", encoding="utf-8") as f:
         yaml.dump(config, f, default_flow_style=None, sort_keys=False, allow_unicode=True, indent=2)
@@ -1733,9 +1858,11 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
                         "The output does not contain any of the keywords."
                     )
 
-        # Verify similarity
+        # Verify similarity (Whisper transcript vs streamed/detokenized text)
         if "text" in modalities and "audio" in modalities:
-            assert response.similarity > 0.9, "The audio content is not same as the text"
+            assert response.similarity is not None and response.similarity > 0.9, (
+                "The audio content is not same as the text"
+            )
             print(f"similarity is: {response.similarity}")
 
 
@@ -2135,7 +2262,11 @@ class OpenAIClientHandler:
         Send OpenAI requests.
 
         Args:
-            request_config: Request configuration dictionary containing parameters like model, messages, stream
+            request_config: Request configuration dictionary containing parameters like model, messages, stream.
+                Optional ``use_audio_in_video`` (bool): when true, sets
+                ``extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}`` for Qwen-Omni video+audio
+                extraction (merged with any existing ``extra_body`` / ``mm_processor_kwargs``).
+                Optional ``extra_body`` (dict): passed through to ``chat.completions.create`` after merge.
             request_num: Number of requests, defaults to 1 (single request)
 
         Returns:
@@ -2146,14 +2277,28 @@ class OpenAIClientHandler:
         stream = request_config.get("stream", False)
         modalities = request_config.get("modalities", ["text", "audio"])
 
+        extra_body: dict[str, Any] = {}
+        raw_extra = request_config.get("extra_body")
+        if raw_extra:
+            extra_body.update(raw_extra)
+        if request_config.get("use_audio_in_video"):
+            mm = dict(extra_body.get("mm_processor_kwargs") or {})
+            mm["use_audio_in_video"] = True
+            extra_body["mm_processor_kwargs"] = mm
+        extra_body_arg: dict[str, Any] | None = extra_body if extra_body else None
+
+        create_kwargs: dict[str, Any] = {
+            "model": request_config.get("model"),
+            "messages": request_config.get("messages"),
+            "stream": stream,
+            "modalities": modalities,
+        }
+        if extra_body_arg is not None:
+            create_kwargs["extra_body"] = extra_body_arg
+
         if request_num == 1:
             # Send single request
-            chat_completion = self.client.chat.completions.create(
-                model=request_config.get("model"),
-                messages=request_config.get("messages"),
-                stream=stream,
-                modalities=modalities,
-            )
+            chat_completion = self.client.chat.completions.create(**create_kwargs)
 
             if stream:
                 response = self._process_stream_omni_response(chat_completion)

tests/diffusion/offloader/test_layerwise_backend.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for LayerwiseOffloadHook."""
+
+import gc
+import os
+import socket
+from contextlib import contextmanager
+
+import pytest
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
+
+import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
+from vllm_omni.diffusion.offloader.layerwise_backend import LayerwiseOffloadHook
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class DummyStream:
+    def wait_stream(self, _stream) -> None:
+        return None
+
+    def wait_event(self, _event) -> None:
+        return None
+
+
+class DummyEvent:
+    def record(self, _stream) -> None:
+        return None
+
+
+@contextmanager
+def dummy_stream(_stream):
+    yield None
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return int(s.getsockname()[1])
+
+
+def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
+    os.environ["RANK"] = str(rank)
+    os.environ["LOCAL_RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(master_port)
+
+
+def _cleanup_distributed() -> None:
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+    for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+        os.environ.pop(key, None)
+
+    gc.collect()
+    if current_omni_platform.is_available():
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()
+
+
+@pytest.fixture(scope="module")
+def dist_group():
+    master_port = _find_free_port()
+    _set_dist_env(rank=0, world_size=1, master_port=master_port)
+
+    dist.init_process_group("gloo", rank=0, world_size=1)
+    try:
+        yield
+    finally:
+        _cleanup_distributed()
+
+
+@pytest.fixture
+def patched_offload_runtime(mocker):
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
+
+
+class TinyBlock(nn.Module):
+    def __init__(self, values: torch.Tensor):
+        super().__init__()
+        mesh = DeviceMesh("cpu", [0])
+        dtensor = DTensor.from_local(values, mesh, [Replicate()])
+        self.weight = nn.Parameter(dtensor)
+
+
+def _make_values(start: float) -> torch.Tensor:
+    return torch.arange(start, start + 4, dtype=torch.float32)
+
+
+class TestLayerwiseOffloadHook:
+    def test_dtensor_wrapper_is_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
+        current_block = TinyBlock(_make_values(1.0))
+        next_block = TinyBlock(_make_values(10.0))
+
+        hook = LayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            stream=DummyStream(),
+            pin_memory=False,
+        )
+
+        hook.initialize_hook(current_block)
+
+        assert isinstance(next_block.weight, DTensor)
+        assert next_block.weight.to_local().is_meta
+        assert next_block.weight.to_local().shape == torch.Size([4])
+        assert hook.dtype_metadata[next_block.weight.dtype][0]["shape"] == torch.Size([4])
+
+        hook.prefetch_layer(non_blocking=False)
+        assert isinstance(next_block.weight, DTensor)
+        assert torch.equal(next_block.weight.to_local(), _make_values(10.0))
+        assert next_block.weight.to_local().shape == torch.Size([4])
+
+        hook.offload_layer()
+        assert isinstance(current_block.weight, DTensor)
+        assert current_block.weight.to_local().is_meta
+        assert current_block.weight.to_local().shape == torch.Size([4])
+        assert not hook.is_materialized

tests/diffusion/test_diffusion_scheduler.py

@@ -1,12 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import queue
 import threading
+from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
 import pytest
+import torch
 
-from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
@@ -14,6 +17,7 @@ from vllm_omni.diffusion.sched import (
     RequestScheduler,
     Scheduler,
     SchedulerInterface,
+    StepScheduler,
 )
 from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -29,9 +33,46 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:
     )
 
 
-def _make_request_output(req_id: str, *, error: str | None = None) -> DiffusionOutput:
-    del req_id
-    return DiffusionOutput(output=None, error=error)
+def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
+    return SimpleNamespace(
+        req_id=req_id,
+        step_index=None,
+        finished=finished,
+        result=DiffusionOutput(output=None, error=error),
+    )
+
+
+def _make_step_output(
+    req_id: str,
+    step_index: int,
+    *,
+    finished: bool = False,
+    error: str | None = None,
+):
+    return SimpleNamespace(
+        req_id=req_id,
+        step_index=step_index,
+        finished=finished,
+        result=DiffusionOutput(output=None, error=error) if error is not None else None,
+    )
+
+
+def _make_step_request(
+    req_id: str,
+    *,
+    num_inference_steps: int = 4,
+    step_index: int | None = None,
+    sampling_params: OmniDiffusionSamplingParams | None = None,
+) -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompts=[f"prompt_{req_id}"],
+        sampling_params=sampling_params
+        or OmniDiffusionSamplingParams(
+            num_inference_steps=num_inference_steps,
+            step_index=step_index,
+        ),
+        request_ids=[req_id],
+    )
 
 
 def _new_ids(sched_output) -> list[str]:
@@ -43,7 +84,7 @@ def _cached_ids(sched_output) -> list[str]:
 
 
 class _StubScheduler(SchedulerInterface):
-    def __init__(self, request: OmniDiffusionRequest, output: DiffusionOutput) -> None:
+    def __init__(self, request: OmniDiffusionRequest, output) -> None:
         self._request = request
         self._output = output
         self.initialized_with = None
@@ -75,9 +116,10 @@ class _StubScheduler(SchedulerInterface):
             is_empty=False,
         )
 
-    def update_from_output(self, sched_output, output: DiffusionOutput) -> set[str]:
+    def update_from_output(self, sched_output, output) -> set[str]:
         del sched_output
         assert output is self._output
+        self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
         return {self._sched_req_id}
 
     def has_requests(self) -> bool:
@@ -185,9 +227,14 @@ class TestRequestScheduler:
         state_b = self.scheduler.get_request_state(req_id_b)
         assert state_b.status == DiffusionRequestStatus.FINISHED_ABORTED
 
+        first = self.scheduler.schedule()
+        assert first.finished_req_ids == {req_id_b}
         # A should still run normally.
-        output_a = self.scheduler.schedule()
-        assert _new_ids(output_a) == [req_id_a]
+        assert _new_ids(first) == [req_id_a]
+
+        # B is already marked finished aborted, scheduling again should not pull it.
+        second = self.scheduler.schedule()
+        assert second.finished_req_ids == set()
 
         # Abort running request.
         self.scheduler.finish_requests(req_id_a, DiffusionRequestStatus.FINISHED_ABORTED)
@@ -233,33 +280,33 @@ class TestDiffusionEngine:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine.scheduler = RequestScheduler()
         engine.scheduler.initialize(Mock())
-        engine.executor = Mock()
-        engine._rpc_lock = threading.Lock()
+        engine._rpc_lock = threading.RLock()
+        engine.abort_queue = queue.Queue()
 
         request = _make_request("engine")
-        expected = DiffusionOutput(output=None)
-        engine.executor.add_req.return_value = expected
+        runner_output = _make_request_output("engine")
+        engine.execute_fn = Mock(return_value=runner_output)
 
         output = engine.add_req_and_wait_for_response(request)
 
-        assert output is expected
-        engine.executor.add_req.assert_called_once_with(request)
+        assert output is runner_output.result
+        engine.execute_fn.assert_called_once()
 
     def test_supports_scheduler_interface_injection(self) -> None:
         request = _make_request("engine_iface")
-        expected = DiffusionOutput(output=None)
-        scheduler = _StubScheduler(request, expected)
+        runner_output = _make_request_output("engine_iface")
+        scheduler = _StubScheduler(request, runner_output)
 
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine.scheduler = scheduler
-        engine.executor = Mock()
-        engine.executor.add_req = Mock(return_value=expected)
-        engine._rpc_lock = threading.Lock()
+        engine._rpc_lock = threading.RLock()
+        engine.abort_queue = queue.Queue()
+        engine.execute_fn = Mock(return_value=runner_output)
 
         output = engine.add_req_and_wait_for_response(request)
 
-        assert output is expected
-        engine.executor.add_req.assert_called_once_with(request)
+        assert output is runner_output.result
+        engine.execute_fn.assert_called_once()
 
     def test_initializes_injected_scheduler(self) -> None:
         request = _make_request("init")
@@ -289,6 +336,59 @@ class TestDiffusionEngine:
         assert req_id in finished
         assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
 
+    def test_step_raises_aborted_error(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.pre_process_func = None
+        engine.add_req_and_wait_for_response = Mock(
+            return_value=DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
+        )
+
+        with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"):
+            engine.step(_make_request("req-abort"))
+
+    def test_abort_queue_marks_request_finished_aborted(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = RequestScheduler()
+        engine.scheduler.initialize(Mock())
+        engine.abort_queue = queue.Queue()
+
+        req_id = engine.scheduler.add_request(_make_request("req-abort"))
+        engine.abort("req-abort")
+        engine._process_aborts_queue()
+
+        assert engine.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_ABORTED
+
+    def test_finalize_finished_request_returns_aborted_output(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = RequestScheduler()
+        engine.scheduler.initialize(Mock())
+
+        req_id = engine.scheduler.add_request(_make_request("req-finalize"))
+        engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
+
+        output = engine._finalize_finished_request(req_id)
+
+        assert output.aborted is True
+        assert output.abort_message == "Request req-finalize aborted."
+
+    def test_initializes_step_scheduler_when_step_execution_enabled(self) -> None:
+        od_config = Mock(model_class_name="mock_model")
+        od_config.step_execution = True
+        fake_executor = Mock()
+        fake_executor_cls = Mock(return_value=fake_executor)
+
+        with (
+            patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", return_value=None),
+            patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", return_value=None),
+            patch("vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", return_value=fake_executor_cls),
+            patch.object(DiffusionEngine, "_dummy_run", return_value=None),
+        ):
+            engine = DiffusionEngine(od_config)
+
+        assert isinstance(engine.scheduler, StepScheduler)
+        assert engine.execute_fn is fake_executor.execute_step
+        fake_executor_cls.assert_called_once_with(od_config)
+
     def test_dummy_run_raises_on_output_error(self) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine.od_config = Mock(model_class_name="mock_model")
@@ -297,3 +397,240 @@ class TestDiffusionEngine:
 
         with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
             engine._dummy_run()
+
+
+class TestStepScheduler:
+    def setup_method(self) -> None:
+        self.scheduler: StepScheduler = StepScheduler()
+        self.scheduler.initialize(Mock())
+
+    def test_single_request_step_lifecycle(self) -> None:
+        request = _make_step_request("step", num_inference_steps=3)
+        req_id = self.scheduler.add_request(request)
+
+        first = self.scheduler.schedule()
+        assert _new_ids(first) == [req_id]
+        assert _cached_ids(first) == []
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 0
+
+        finished = self.scheduler.update_from_output(first, _make_step_output(req_id, step_index=1))
+        assert finished == set()
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING
+        assert request.sampling_params.step_index == 1
+        assert self.scheduler.has_requests() is True
+
+        second = self.scheduler.schedule()
+        assert _new_ids(second) == []
+        assert _cached_ids(second) == [req_id]
+        assert second.num_running_reqs == 1
+        assert second.num_waiting_reqs == 0
+
+        finished = self.scheduler.update_from_output(second, _make_step_output(req_id, step_index=2))
+        assert finished == set()
+        assert request.sampling_params.step_index == 2
+
+        third = self.scheduler.schedule()
+        assert _new_ids(third) == []
+        assert _cached_ids(third) == [req_id]
+
+        finished = self.scheduler.update_from_output(
+            third,
+            _make_step_output(req_id, step_index=3, finished=True),
+        )
+        assert finished == {req_id}
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
+        assert request.sampling_params.step_index == 3
+        assert self.scheduler.has_requests() is False
+
+    def test_fifo_single_request_scheduling(self) -> None:
+        req_id_a = self.scheduler.add_request(_make_step_request("a", num_inference_steps=2))
+        req_id_b = self.scheduler.add_request(_make_step_request("b", num_inference_steps=2))
+
+        first = self.scheduler.schedule()
+        assert _new_ids(first) == [req_id_a]
+        assert _cached_ids(first) == []
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
+        finished = self.scheduler.update_from_output(first, _make_step_output(req_id_a, step_index=1))
+        assert finished == set()
+
+        second = self.scheduler.schedule()
+        assert _new_ids(second) == []
+        assert _cached_ids(second) == [req_id_a]
+        assert second.num_running_reqs == 1
+        assert second.num_waiting_reqs == 1
+
+        finished = self.scheduler.update_from_output(
+            second,
+            _make_step_output(req_id_a, step_index=2, finished=True),
+        )
+        assert finished == {req_id_a}
+
+        third = self.scheduler.schedule()
+        assert _new_ids(third) == [req_id_b]
+        assert _cached_ids(third) == []
+        assert third.num_running_reqs == 1
+        assert third.num_waiting_reqs == 0
+
+    def test_error_output_marks_finished_error(self) -> None:
+        req_id = self.scheduler.add_request(_make_step_request("err", num_inference_steps=3))
+
+        sched_output = self.scheduler.schedule()
+        assert _new_ids(sched_output) == [req_id]
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            _make_step_output(req_id, step_index=1, finished=True, error="worker failed"),
+        )
+
+        assert finished == {req_id}
+        state = self.scheduler.get_request_state(req_id)
+        assert state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert state.error == "worker failed"
+        assert self.scheduler.has_requests() is False
+
+    def test_missing_step_index_marks_finished_error(self) -> None:
+        req_id = self.scheduler.add_request(_make_step_request("missing", num_inference_steps=3))
+
+        sched_output = self.scheduler.schedule()
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            SimpleNamespace(
+                req_id=req_id,
+                step_index=None,
+                finished=True,
+                result=None,
+            ),
+        )
+
+        assert finished == {req_id}
+        state = self.scheduler.get_request_state(req_id)
+        assert state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert state.error == "Missing step_index in RunnerOutput"
+
+    def test_abort_request_for_waiting_and_running(self) -> None:
+        req_id_a = self.scheduler.add_request(_make_step_request("a", num_inference_steps=2))
+        req_id_b = self.scheduler.add_request(_make_step_request("b", num_inference_steps=2))
+
+        self.scheduler.finish_requests(req_id_b, DiffusionRequestStatus.FINISHED_ABORTED)
+        assert self.scheduler.get_request_state(req_id_b).status == DiffusionRequestStatus.FINISHED_ABORTED
+
+        running = self.scheduler.schedule()
+        assert _new_ids(running) == [req_id_a]
+
+        self.scheduler.finish_requests(req_id_a, DiffusionRequestStatus.FINISHED_ABORTED)
+        assert self.scheduler.get_request_state(req_id_a).status == DiffusionRequestStatus.FINISHED_ABORTED
+        assert self.scheduler.has_requests() is False
+
+    def test_has_requests_state_transition(self) -> None:
+        assert self.scheduler.has_requests() is False
+
+        req_id = self.scheduler.add_request(_make_step_request("has", num_inference_steps=2))
+        assert self.scheduler.has_requests() is True
+
+        sched_output = self.scheduler.schedule()
+        assert self.scheduler.has_requests() is True
+
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            _make_step_output(req_id, step_index=2, finished=True),
+        )
+        assert finished == {req_id}
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
+        assert self.scheduler.has_requests() is False
+
+    def test_scheduled_request_aborted_before_update_is_returned_finished(self) -> None:
+        req_id = self.scheduler.add_request(_make_step_request("abort-late", num_inference_steps=2))
+
+        sched_output = self.scheduler.schedule()
+        self.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
+
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            _make_step_output(req_id, step_index=1),
+        )
+        assert finished == {req_id}
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_ABORTED
+
+    def test_preempt_request_preserves_step_index(self) -> None:
+        request = _make_step_request("preempt", num_inference_steps=3)
+        req_id = self.scheduler.add_request(request)
+
+        first = self.scheduler.schedule()
+        assert self.scheduler.update_from_output(first, _make_step_output(req_id, step_index=1)) == set()
+        assert request.sampling_params.step_index == 1
+
+        second = self.scheduler.schedule()
+        assert _cached_ids(second) == [req_id]
+        assert self.scheduler.preempt_request(req_id) is True
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.PREEMPTED
+        assert request.sampling_params.step_index == 1
+
+        third = self.scheduler.schedule()
+        assert _cached_ids(third) == [req_id]
+        assert request.sampling_params.step_index == 1
+
+    @pytest.mark.parametrize(
+        ("sampling_params", "expected_steps"),
+        [
+            (
+                OmniDiffusionSamplingParams(
+                    timesteps=torch.tensor([1.0, 0.5, 0.0]),
+                    sigmas=[1.0, 0.5, 0.25, 0.0],
+                    num_inference_steps=5,
+                ),
+                3,
+            ),
+            (
+                OmniDiffusionSamplingParams(
+                    sigmas=[1.0, 0.5],
+                    num_inference_steps=5,
+                ),
+                2,
+            ),
+            (
+                OmniDiffusionSamplingParams(
+                    num_inference_steps=4,
+                ),
+                4,
+            ),
+        ],
+    )
+    def test_total_steps_priority(self, sampling_params: OmniDiffusionSamplingParams, expected_steps: int) -> None:
+        request = _make_step_request("priority", sampling_params=sampling_params)
+        req_id = self.scheduler.add_request(request)
+
+        for _ in range(expected_steps - 1):
+            sched_output = self.scheduler.schedule()
+            assert sched_output.scheduled_req_ids == [req_id]
+            next_step = request.sampling_params.step_index + 1
+            assert (
+                self.scheduler.update_from_output(
+                    sched_output,
+                    _make_step_output(req_id, step_index=next_step),
+                )
+                == set()
+            )
+
+        final_output = self.scheduler.schedule()
+        assert final_output.scheduled_req_ids == [req_id]
+        assert self.scheduler.update_from_output(
+            final_output,
+            _make_step_output(req_id, step_index=expected_steps, finished=True),
+        ) == {req_id}
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
+
+    @pytest.mark.parametrize(
+        "sampling_params",
+        [
+            OmniDiffusionSamplingParams(num_inference_steps=0),
+            OmniDiffusionSamplingParams(num_inference_steps=3, step_index=3),
+            OmniDiffusionSamplingParams(num_inference_steps=3, step_index=-1),
+        ],
+    )
+    def test_rejects_invalid_initial_step_state(self, sampling_params: OmniDiffusionSamplingParams) -> None:
+        request = _make_step_request("invalid", sampling_params=sampling_params)
+
+        with pytest.raises(ValueError):
+            self.scheduler.add_request(request)

tests/diffusion/test_diffusion_step_pipeline.py

@@ -1,10 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for step-level diffusion runner and worker execution."""
+"""Tests for step-level diffusion execution across runner / worker / executor / engine."""
 
 import os
+import queue
+import threading
 from contextlib import contextmanager
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -12,6 +15,7 @@ import torch
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
 from tests.utils import hardware_test
 from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -20,10 +24,13 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
+from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.ipc import (
     pack_diffusion_output_shm,
     unpack_diffusion_output_shm,
 )
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched import StepScheduler
 from vllm_omni.diffusion.sched.interface import (
     CachedRequestData,
     DiffusionSchedulerOutput,
@@ -32,6 +39,8 @@ from vllm_omni.diffusion.sched.interface import (
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 from vllm_omni.diffusion.worker.utils import RunnerOutput
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@@ -86,6 +95,23 @@ class _StepPipeline:
         return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
 
 
+class _InterruptingStepPipeline(_StepPipeline):
+    interrupt = True
+
+    def denoise_step(self, state, **kwargs):
+        del state, kwargs
+        self.denoise_calls += 1
+        return None
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        del state, noise_pred, kwargs
+        raise AssertionError("step_scheduler should not run after interrupt")
+
+    def post_decode(self, state, **kwargs):
+        del state, kwargs
+        raise AssertionError("post_decode should not run after interrupt")
+
+
 class _IdentityNoiseTransformer(torch.nn.Module):
     def forward(self, x: torch.Tensor, **kwargs):
         del kwargs
@@ -188,6 +214,21 @@ def _make_step_request(num_inference_steps: int = 2):
     )
 
 
+def _assert_aborted_output(output: DiffusionOutput, request_id: str) -> None:
+    assert output.output is None
+    assert output.error is None
+    assert output.aborted is True
+    assert output.abort_message == f"Request {request_id} aborted."
+
+
+def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompts=[f"prompt-{req_id}"],
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
+        request_ids=[req_id],
+    )
+
+
 def _make_runner():
     runner = object.__new__(DiffusionModelRunner)
     runner.vllm_config = object()
@@ -242,6 +283,18 @@ def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_
     )
 
 
+def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine:
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline")
+    engine.pre_process_func = None
+    engine.post_process_func = None
+    engine.scheduler = scheduler
+    engine.execute_fn = execute_fn
+    engine._rpc_lock = threading.RLock()
+    engine.abort_queue = queue.Queue()
+    return engine
+
+
 def _expected_output_for_mode(mode: str) -> torch.Tensor:
     if mode == "cfg":
         return torch.tensor([[3.0]])
@@ -322,6 +375,52 @@ class TestRunner:
         assert runner.pipeline.scheduler_calls == 2
         assert runner.pipeline.decode_calls == 1
 
+    def test_rejects_multi_request_step_batch(self):
+        runner = _make_runner()
+        req_1 = _make_step_request()
+        req_2 = _make_step_request()
+        req_2.request_ids = ["req-2"]
+
+        scheduler_output = DiffusionSchedulerOutput(
+            step_id=0,
+            scheduled_new_reqs=[
+                NewRequestData(sched_req_id="req-1", req=req_1),
+                NewRequestData(sched_req_id="req-2", req=req_2),
+            ],
+            scheduled_cached_reqs=CachedRequestData.make_empty(),
+            finished_req_ids=set(),
+            num_running_reqs=2,
+            num_waiting_reqs=0,
+        )
+
+        with pytest.raises(ValueError, match="batch_size=1"):
+            DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+
+    def test_rejects_missing_cached_state(self):
+        runner = _make_runner()
+
+        with pytest.raises(ValueError, match="Missing cached state"):
+            DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(sched_req_id="req-missing"))
+
+    def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch):
+        runner = _make_runner()
+        runner.pipeline = _InterruptingStepPipeline()
+        req = _make_step_request()
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+        output = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
+
+        assert output.req_id == "req-1"
+        assert output.step_index == 0
+        assert output.finished is True
+        assert output.result is not None
+        assert output.result.error == "stepwise denoise interrupted"
+        assert "req-1" not in runner.state_cache
+        assert runner.pipeline.prepare_calls == 1
+        assert runner.pipeline.denoise_calls == 1
+        assert runner.pipeline.scheduler_calls == 0
+        assert runner.pipeline.decode_calls == 0
+
     def test_load_model_rejects_unsupported_step_execution(self, monkeypatch):
         class _RequestOnlyPipeline:
             pass
@@ -439,6 +538,153 @@ class TestWorker:
             DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
 
+@pytest.mark.cpu
+class TestExecutor:
+    """MultiprocDiffusionExecutor.execute_step"""
+
+    def test_execute_step_passes_through_runner_output(self):
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._ensure_open = lambda: None
+        expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None)
+        executor.collective_rpc = Mock(return_value=expected)
+
+        request = _make_engine_request("req-step", num_inference_steps=2)
+        scheduler_output = _make_scheduler_output(request, sched_req_id="req-step")
+
+        output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output)
+
+        assert output is expected
+
+
+@pytest.mark.cpu
+class TestEngine:
+    """Step-execution paths in DiffusionEngine.add_req_and_wait_for_response"""
+
+    @pytest.mark.parametrize(
+        ("execute_fn", "expected_error"),
+        [
+            (
+                lambda _: RunnerOutput(
+                    req_id="req-error",
+                    step_index=1,
+                    finished=True,
+                    result=DiffusionOutput(error="boom"),
+                ),
+                "boom",
+            ),
+            (
+                lambda _: (_ for _ in ()).throw(RuntimeError("gpu on fire")),
+                "gpu on fire",
+            ),
+        ],
+    )
+    def test_step_engine_returns_error(self, execute_fn, expected_error):
+        scheduler = StepScheduler()
+        scheduler.initialize(Mock())
+        engine = _make_engine(scheduler, execute_fn=execute_fn)
+
+        output = engine.add_req_and_wait_for_response(_make_engine_request("req-error", num_inference_steps=2))
+
+        assert output.output is None
+        assert expected_error in output.error
+
+    def test_step_execution_completes(self):
+        scheduler = StepScheduler()
+        scheduler.initialize(Mock())
+        engine = _make_engine(scheduler)
+        request = _make_engine_request("req-step", num_inference_steps=2)
+
+        call_count = {"n": 0}
+
+        def execute_fn(_):
+            call_count["n"] += 1
+            finished = call_count["n"] == 2
+            return RunnerOutput(
+                req_id="req-step",
+                step_index=call_count["n"],
+                finished=finished,
+                result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None),
+            )
+
+        engine.execute_fn = execute_fn
+
+        output = engine.add_req_and_wait_for_response(request)
+
+        assert call_count["n"] == 2
+        assert output.error is None
+        assert torch.equal(output.output, torch.tensor([2.0]))
+
+    def test_step_abort_stops_rescheduling_after_first_step(self):
+        scheduler = StepScheduler()
+        scheduler.initialize(Mock())
+        engine = _make_engine(scheduler)
+        request = _make_engine_request("req-stop", num_inference_steps=4)
+
+        step = {"n": 0}
+
+        def execute_fn(_):
+            step["n"] += 1
+            engine.abort("req-stop")
+            return RunnerOutput(
+                req_id="req-stop",
+                step_index=1,
+                finished=False,
+                result=None,
+            )
+
+        engine.execute_fn = execute_fn
+
+        output = engine.add_req_and_wait_for_response(request)
+
+        assert step["n"] == 1
+        _assert_aborted_output(output, "req-stop")
+
+    def test_step_abort_after_reschedule_returns_aborted_output(self):
+        scheduler = StepScheduler()
+        scheduler.initialize(Mock())
+        engine = _make_engine(scheduler)
+        request = _make_engine_request("req-mid", num_inference_steps=4)
+
+        step = {"n": 0}
+
+        def execute_fn(sched_output):
+            step["n"] += 1
+            if step["n"] == 2:
+                assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1)
+                engine.abort("req-mid")
+            return RunnerOutput(
+                req_id="req-mid",
+                step_index=step["n"],
+                finished=False,
+                result=None,
+            )
+
+        engine.execute_fn = execute_fn
+
+        output = engine.add_req_and_wait_for_response(request)
+
+        assert step["n"] == 2
+        _assert_aborted_output(output, "req-mid")
+
+    def test_finished_step_without_result_returns_error(self):
+        scheduler = StepScheduler()
+        scheduler.initialize(Mock())
+        engine = _make_engine(
+            scheduler,
+            execute_fn=lambda _: RunnerOutput(
+                req_id="req-missing",
+                step_index=1,
+                finished=True,
+                result=None,
+            ),
+        )
+
+        output = engine.add_req_and_wait_for_response(_make_engine_request("req-missing", num_inference_steps=1))
+
+        assert output.output is None
+        assert output.error == "Diffusion execution finished without a final output."
+
+
 @pytest.mark.cpu
 class TestIPC:
     def test_pack_unpack_runner_output_shm(self):
@@ -458,6 +704,15 @@ class TestIPC:
 class TestSupportedPipelines:
     """Step-execution protocol checks for supported pipelines."""
 
+    def test_default_stage_config_includes_step_execution(self):
+        stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+            {
+                "step_execution": True,
+            }
+        )[0]
+
+        assert stage_cfg["engine_args"]["step_execution"] is True
+
     def test_qwen_image_supports_step_execution(self):
         from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
         from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -66,7 +66,9 @@ def _make_engine(num_gpus: int = 1):
     sched.initialize(Mock())
     engine.scheduler = sched
     engine.executor = executor
-    engine._rpc_lock = threading.Lock()
+    engine._rpc_lock = threading.RLock()
+    engine.abort_queue = queue.Queue()
+    engine.execute_fn = executor.execute_request
     return engine, executor, req_q, res_q
 
 
@@ -80,7 +82,7 @@ def _start_worker(req_q, res_q, count=2):
             req = req_q.get(timeout=10)
             method = req.get("method", "")
             args = req.get("args", ())
-            if method == "generate" and args and hasattr(args[0], "request_ids"):
+            if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_ids"):
                 tag = f"result_for_{args[0].request_ids[0]}"
             elif args:
                 tag = f"result_for_{args[0]}"
@@ -116,11 +118,11 @@ def _inject_interleave(executor):
     return a_enqueued, b_complete
 
 
-# ──────────────────── bug-reproduction: concurrent add_req ────────────────
+# ───────────────── concurrent request execution ─────────────────
 
 
-class TestConcurrentAddReqBug:
-    """Two concurrent ``add_req_and_wait_for_response()`` calls swap results."""
+class TestConcurrentRequestExecution:
+    """Concurrent request execution should not swap results."""
 
     def test_results_are_correctly_routed(self):
         engine, executor, req_q, res_q = _make_engine()
@@ -151,11 +153,11 @@ class TestConcurrentAddReqBug:
         assert results["B"].error == "result_for_B"
 
 
-# ──────────────── bug-reproduction: concurrent collective_rpc ─────────────
+# ───────────────── concurrent collective RPC ─────────────────
 
 
-class TestConcurrentCollectiveRpcBug:
-    """Two concurrent ``collective_rpc()`` calls swap results."""
+class TestConcurrentCollectiveRpc:
+    """Concurrent ``collective_rpc()`` calls should not swap results."""
 
     def test_results_are_correctly_routed(self):
         engine, executor, req_q, res_q = _make_engine()
@@ -192,11 +194,11 @@ class TestConcurrentCollectiveRpcBug:
         assert results["B"].error == "result_for_call_B"
 
 
-# ──────── bug-reproduction: add_req vs collective_rpc concurrently ────────
+# ──────────── concurrent request execution and collective RPC ────────────
 
 
-class TestConcurrentAddReqVsCollectiveRpcBug:
-    """``add_req`` and ``collective_rpc`` running concurrently swap results."""
+class TestConcurrentRequestExecutionAndCollectiveRpc:
+    """Request execution and ``collective_rpc()`` should not swap results."""
 
     def test_results_are_correctly_routed(self):
         engine, executor, req_q, res_q = _make_engine()
@@ -205,7 +207,7 @@ class TestConcurrentAddReqVsCollectiveRpcBug:
 
         results: dict[str, object] = {}
 
-        def _a():  # add_req path
+        def _a():  # request execution path
             results["A"] = engine.add_req_and_wait_for_response(_mock_request("A"))
 
         def _b():  # collective_rpc path
@@ -230,10 +232,10 @@ class TestConcurrentAddReqVsCollectiveRpcBug:
         assert results["B"].error == "result_for_call_B"
 
 
-# ─────────────── backward-compatibility (serial) tests ────────────────────
+# ─────────────────────── serial operation coverage ───────────────────────
 
 
-class TestSerialOperations:
+class TestSerialEngineOperations:
     """Verify correct behaviour for single-threaded (serial) usage.
 
     These tests must pass both **before** and **after** any concurrency fix
@@ -385,18 +387,18 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
 
         executor._result_mq.dequeue = _hanging_dequeue
 
-        # Thread running add_req — acquires the lock, enqueues, then
+        # Thread running request execution — acquires the lock, enqueues, then
         # blocks on dequeue forever (worker hang).
-        def _stalled_add_req():
+        def _stalled_request_execution():
             try:
                 engine.add_req_and_wait_for_response(_mock_request("stalled"))
             except Exception:
                 pass
 
-        t = threading.Thread(target=_stalled_add_req, daemon=True)
+        t = threading.Thread(target=_stalled_request_execution, daemon=True)
         t.start()
 
-        # Wait until add_req is truly inside the lock and blocking.
+        # Wait until request execution is truly inside the lock and blocking.
         add_req_blocked.wait(5)
 
         # collective_rpc should time out at lock acquisition, not hang.

tests/e2e/accuracy/conftest.py

@@ -47,6 +47,19 @@ def pytest_addoption(parser):
         help="Balanced sample count per GEdit task group",
     )
     group.addoption("--accuracy-workers", action="store", type=int, default=1, help="Worker count for accuracy benches")
+    group.addoption(
+        "--wan22-i2v-image-source",
+        action="store",
+        default=None,
+        help="Image source for Wan2.2 I2V accuracy tests. Can be local path or remote URL.",
+    )
+    group.addoption(
+        "--wan22-i2v-online-timeout-seconds",
+        action="store",
+        type=int,
+        default=1200,
+        help="Online serving timeout in seconds for Wan2.2 I2V accuracy tests.",
+    )
 
 
 def _hf_cache_root() -> Path:
@@ -142,6 +155,17 @@ def accuracy_workers(request: pytest.FixtureRequest) -> int:
     return int(request.config.getoption("accuracy_workers"))
 
 
+@pytest.fixture(scope="session")
+def wan22_i2v_image_source(request: pytest.FixtureRequest) -> str | None:
+    value = request.config.getoption("wan22_i2v_image_source")
+    return str(value) if value else None
+
+
+@pytest.fixture(scope="session")
+def wan22_i2v_online_timeout_seconds(request: pytest.FixtureRequest) -> int:
+    return int(request.config.getoption("wan22_i2v_online_timeout_seconds"))
+
+
 @pytest.fixture(scope="session")
 def gebench_samples_per_type(request: pytest.FixtureRequest) -> int:
     return int(request.config.getoption("gebench_samples_per_type"))

tests/e2e/accuracy/wan22_i2v/run_wan22_i2v_diffusers_cp.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import argparse
+import base64
+import json
+from io import BytesIO
+from pathlib import Path
+
+import requests
+import torch
+from diffusers import UniPCMultistepScheduler, WanImageToVideoPipeline
+from diffusers.pipelines.wan import pipeline_wan_i2v as wan_i2v_module
+from diffusers.utils import export_to_video, load_image
+from PIL import Image
+
+from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import BOUNDARY_RATIO
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run Wan2.2 I2V diffusers offline generation.")
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--image-source", required=True)
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--negative-prompt", required=True)
+    parser.add_argument("--size", required=True)
+    parser.add_argument("--fps", type=int, required=True)
+    parser.add_argument("--num-frames", type=int, required=True)
+    parser.add_argument("--guidance-scale", type=float, required=True)
+    parser.add_argument("--guidance-scale-2", type=float, required=True)
+    parser.add_argument("--flow-shift", type=float, required=True)
+    parser.add_argument("--num-inference-steps", type=int, required=True)
+    parser.add_argument("--seed", type=int, required=True)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--metadata-output", required=True)
+    return parser.parse_args()
+
+
+def _parse_size(size: str) -> tuple[int, int]:
+    width_str, height_str = size.lower().split("x", 1)
+    return int(width_str), int(height_str)
+
+
+class _IdentityFtfy:
+    @staticmethod
+    def fix_text(text: str) -> str:
+        return text
+
+
+def _ensure_wan_ftfy_fallback() -> None:
+    if not hasattr(wan_i2v_module, "ftfy"):
+        wan_i2v_module.ftfy = _IdentityFtfy()
+
+
+def _offline_cuda_device() -> torch.device:
+    return torch.device("cuda:0")
+
+
+def _load_input_image(source: str) -> Image.Image:
+    if source.startswith("data:image"):
+        _, encoded = source.split(",", 1)
+        image = Image.open(BytesIO(base64.b64decode(encoded)))
+        image.load()
+        return image.convert("RGB")
+
+    source_path = Path(source)
+    if source_path.exists():
+        image = Image.open(source_path)
+        image.load()
+        return image.convert("RGB")
+
+    image = load_image(source)
+    if isinstance(image, Image.Image):
+        image.load()
+        return image.convert("RGB")
+
+    response = requests.get(source, timeout=60)
+    response.raise_for_status()
+    image = Image.open(BytesIO(response.content))
+    image.load()
+    return image.convert("RGB")
+
+
+def _resize_to_target(image: Image.Image, *, width: int, height: int) -> Image.Image:
+    return image.resize((width, height), Image.Resampling.LANCZOS)
+
+
+def _configure_scheduler(pipe: WanImageToVideoPipeline, *, flow_shift: float) -> None:
+    pipe.scheduler = UniPCMultistepScheduler.from_config(
+        pipe.scheduler.config,
+        flow_shift=flow_shift,
+    )
+
+
+def _write_metadata(
+    path: Path,
+    *,
+    args: argparse.Namespace,
+    width: int,
+    height: int,
+    frame_count: int,
+) -> None:
+    payload = {
+        "model": args.model,
+        "image_source": args.image_source,
+        "size": args.size,
+        "width": width,
+        "height": height,
+        "fps": args.fps,
+        "num_frames": args.num_frames,
+        "actual_frame_count": frame_count,
+        "guidance_scale": args.guidance_scale,
+        "guidance_scale_2": args.guidance_scale_2,
+        "boundary_ratio": BOUNDARY_RATIO,
+        "flow_shift": args.flow_shift,
+        "num_inference_steps": args.num_inference_steps,
+        "seed": args.seed,
+        "world_size": 1,
+    }
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
+
+
+def main() -> int:
+    args = _parse_args()
+    device = _offline_cuda_device()
+    torch.cuda.set_device(device)
+    _ensure_wan_ftfy_fallback()
+
+    pipe = WanImageToVideoPipeline.from_pretrained(args.model, torch_dtype=torch.bfloat16)
+    pipe.register_to_config(boundary_ratio=BOUNDARY_RATIO)
+    _configure_scheduler(pipe, flow_shift=args.flow_shift)
+    pipe.to(device)
+    pipe.set_progress_bar_config(disable=False)
+
+    input_image = _load_input_image(args.image_source)
+    width, height = _parse_size(args.size)
+    resized_image = _resize_to_target(input_image, width=width, height=height)
+
+    generator = torch.Generator(device=device.type).manual_seed(args.seed)
+    frames = pipe(
+        image=resized_image,
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        height=height,
+        width=width,
+        num_frames=args.num_frames,
+        guidance_scale=args.guidance_scale,
+        guidance_scale_2=args.guidance_scale_2,
+        num_inference_steps=args.num_inference_steps,
+        generator=generator,
+    ).frames[0]
+
+    output_path = Path(args.output)
+    metadata_path = Path(args.metadata_output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    export_to_video(frames, str(output_path), fps=args.fps)
+    _write_metadata(metadata_path, args=args, width=width, height=height, frame_count=len(frames))
+
+    if hasattr(pipe, "maybe_free_model_hooks"):
+        pipe.maybe_free_model_hooks()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py

@@ -0,0 +1,672 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+import time
+from base64 import b64decode, b64encode
+from fractions import Fraction
+from hashlib import sha1
+from io import BytesIO
+from mimetypes import guess_type
+from pathlib import Path
+from typing import Any
+from urllib.parse import urlparse
+
+import pytest
+import requests
+import torch
+from diffusers import UniPCMultistepScheduler
+from PIL import Image
+
+from tests.conftest import OmniServerParams
+from tests.e2e.accuracy.wan22_i2v.run_wan22_i2v_diffusers_cp import (
+    _configure_scheduler,
+    _ensure_wan_ftfy_fallback,
+    _IdentityFtfy,
+    _offline_cuda_device,
+    _resize_to_target,
+)
+from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import (
+    FLOW_SHIFT,
+    FPS,
+    GUIDANCE_SCALE,
+    GUIDANCE_SCALE_2,
+    HEIGHT,
+    MODEL_NAME,
+    NEGATIVE_PROMPT,
+    NUM_FRAMES,
+    NUM_INFERENCE_STEPS,
+    PROMPT,
+    PSNR_THRESHOLD,
+    RABBIT_IMAGE_URL,
+    SEED,
+    SIZE,
+    SSIM_THRESHOLD,
+    WIDTH,
+)
+from tests.utils import hardware_test
+
+
+def test_parse_video_metadata_extracts_dimensions_and_fps() -> None:
+    payload = {
+        "streams": [
+            {
+                "width": 832,
+                "height": 480,
+                "avg_frame_rate": "16/1",
+                "nb_read_frames": "93",
+            }
+        ]
+    }
+
+    metadata = _parse_video_metadata(payload)
+
+    assert metadata["width"] == 832
+    assert metadata["height"] == 480
+    assert metadata["fps"] == 16.0
+    assert metadata["frame_count"] == 93
+
+
+def test_parse_ssim_summary_extracts_all_score() -> None:
+    output = """
+    [Parsed_ssim_0 @ 000001] SSIM Y:0.971903 (15.512007) U:0.965077 (14.569044) V:0.962414 (14.252637) All:0.968311 (15.035654)
+    """
+
+    assert _parse_ssim_score(output) == 0.968311
+
+
+def test_parse_psnr_summary_extracts_average_score() -> None:
+    output = """
+    [Parsed_psnr_0 @ 000001] PSNR y:32.670157 u:31.844621 v:31.513839 average:32.148004 min:31.563744 max:33.201457
+    """
+
+    assert _parse_psnr_score(output) == 32.148004
+
+
+def test_build_diffusers_command_uses_python_runner_path(tmp_path: Path) -> None:
+    runner_path = tmp_path / "run_wan22_i2v_diffusers_cp.py"
+    command = _build_diffusers_command(
+        runner_path=runner_path,
+        image_source=RABBIT_IMAGE_URL,
+        output_path=tmp_path / "offline.mp4",
+        metadata_path=tmp_path / "offline.json",
+    )
+
+    assert command[:2] == [
+        sys.executable,
+        str(runner_path),
+    ]
+    assert "--output" in command
+    assert "--metadata-output" in command
+
+
+def test_resolve_image_source_prefers_existing_local_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    image_path = tmp_path / "rabbit.png"
+    Image.new("RGB", (8, 8), color=(255, 128, 64)).save(image_path)
+
+    assert _resolve_image_source(str(image_path)) == str(image_path.resolve())
+
+
+def test_build_online_image_reference_uses_data_url_for_local_path(tmp_path: Path) -> None:
+    image_path = tmp_path / "rabbit.png"
+    Image.new("RGB", (4, 2), color=(10, 20, 30)).save(image_path)
+
+    reference = _build_online_image_reference(str(image_path))
+
+    assert reference.startswith("data:image/png;base64,")
+    encoded = reference.split(",", 1)[1]
+    assert len(b64decode(encoded)) > 0
+
+
+def test_send_video_request_with_timeout_uses_requested_timeout(monkeypatch: pytest.MonkeyPatch) -> None:
+    class DummyResponse:
+        def raise_for_status(self) -> None:
+            return None
+
+        def json(self) -> dict[str, str]:
+            return {"id": "video_123"}
+
+    class DummyClient:
+        def __init__(self) -> None:
+            self.run_level = "core_model"
+            self.wait_args: tuple[str, int, int] | None = None
+
+        def _build_url(self, path: str) -> str:
+            return f"http://localhost:8000{path}"
+
+        def _wait_until_video_completed(
+            self,
+            video_id: str,
+            poll_interval_seconds: int = 2,
+            timeout_seconds: int = 300,
+        ) -> None:
+            self.wait_args = (video_id, poll_interval_seconds, timeout_seconds)
+
+        def _download_video_content(self, video_id: str) -> bytes:
+            return b"video-bytes"
+
+    monkeypatch.setattr(requests, "post", lambda *args, **kwargs: DummyResponse())
+
+    client = DummyClient()
+    result = _send_video_request_with_timeout(
+        client,
+        {
+            "form_data": {"prompt": "test"},
+            "image_reference": RABBIT_IMAGE_URL,
+        },
+        timeout_seconds=1200,
+    )
+
+    assert result == b"video-bytes"
+    assert client.wait_args == ("video_123", 2, 1200)
+
+
+def test_online_timeout_defaults_to_1200() -> None:
+    assert _online_timeout_seconds(None) == 1200
+
+
+def test_artifact_dir_is_under_repo_result_folder(tmp_path: Path) -> None:
+    artifact_dir = _artifact_dir(str(tmp_path / "rabbit.png"))
+
+    assert artifact_dir.parent == Path(__file__).parent / "result"
+    assert artifact_dir.name.startswith("rabbit-")
+
+
+def test_offline_cuda_device_uses_indexed_cuda_device() -> None:
+    assert _offline_cuda_device() == torch.device("cuda:0")
+
+
+def test_ensure_wan_ftfy_fallback_sets_identity(monkeypatch: pytest.MonkeyPatch) -> None:
+    from diffusers.pipelines.wan import pipeline_wan_i2v as wan_i2v_module
+
+    monkeypatch.delattr(wan_i2v_module, "ftfy", raising=False)
+    _ensure_wan_ftfy_fallback()
+
+    assert hasattr(wan_i2v_module, "ftfy")
+    assert isinstance(wan_i2v_module.ftfy, _IdentityFtfy)
+    assert wan_i2v_module.ftfy.fix_text("abc") == "abc"
+
+
+def test_resize_to_target_matches_requested_dimensions() -> None:
+    image = Image.new("RGB", (100, 50), color=(1, 2, 3))
+
+    resized = _resize_to_target(image, width=832, height=480)
+
+    assert resized.size == (832, 480)
+
+
+def test_configure_scheduler_switches_to_unipc() -> None:
+    class DummyPipe:
+        def __init__(self) -> None:
+            self.scheduler = UniPCMultistepScheduler(num_train_timesteps=1000, flow_shift=1.0)
+
+    pipe = DummyPipe()
+    _configure_scheduler(pipe, flow_shift=5.0)
+
+    assert isinstance(pipe.scheduler, UniPCMultistepScheduler)
+    assert pipe.scheduler.config.flow_shift == 5.0
+
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+WORKSPACE_ROOT = REPO_ROOT.parent
+RUNNER_PATH = Path(__file__).with_name("run_wan22_i2v_diffusers_cp.py")
+RESULT_ROOT = Path(__file__).parent / "result"
+VIDEO_TIMEOUT_SECONDS = 60 * 60
+_SSIM_RE = re.compile(r"All:(?P<score>[0-9.]+)")
+_PSNR_RE = re.compile(r"average:(?P<score>[0-9.]+)")
+SERVER_CASES = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL_NAME,
+            server_args=[
+                "--usp",
+                "2",
+                "--use-hsdp",
+                "--hsdp-shard-size",
+                "2",
+            ],
+            use_omni=True,
+        ),
+        id="wan22_i2v_usp2_hsdp2",
+    )
+]
+
+
+def _parse_video_metadata(payload: dict[str, Any]) -> dict[str, int | float]:
+    streams = payload.get("streams")
+    if not isinstance(streams, list) or not streams:
+        raise ValueError(f"ffprobe payload did not include video streams: {payload}")
+
+    stream = streams[0]
+    width = int(stream["width"])
+    height = int(stream["height"])
+    fps = float(Fraction(stream["avg_frame_rate"]))
+    frame_count_value = stream.get("nb_read_frames") or stream.get("nb_frames")
+    if frame_count_value is None:
+        raise ValueError(f"ffprobe payload did not include frame count: {payload}")
+    frame_count = int(frame_count_value)
+    return {
+        "width": width,
+        "height": height,
+        "fps": fps,
+        "frame_count": frame_count,
+    }
+
+
+def _parse_ssim_score(output: str) -> float:
+    match = _SSIM_RE.search(output)
+    if match is None:
+        raise ValueError(f"Could not parse SSIM score from ffmpeg output:\n{output}")
+    return float(match.group("score"))
+
+
+def _parse_psnr_score(output: str) -> float:
+    match = _PSNR_RE.search(output)
+    if match is None:
+        raise ValueError(f"Could not parse PSNR score from ffmpeg output:\n{output}")
+    return float(match.group("score"))
+
+
+def _probe_binary(binary: str) -> str:
+    resolved = shutil.which(binary)
+    if resolved is None:
+        pytest.skip(f"{binary} is required for Wan2.2 video similarity e2e test.")
+    return resolved
+
+
+def _build_diffusers_command(
+    *,
+    runner_path: Path,
+    image_source: str,
+    output_path: Path,
+    metadata_path: Path,
+) -> list[str]:
+    return [
+        sys.executable,
+        str(runner_path),
+        "--model",
+        MODEL_NAME,
+        "--image-source",
+        image_source,
+        "--prompt",
+        PROMPT,
+        "--negative-prompt",
+        NEGATIVE_PROMPT,
+        "--size",
+        SIZE,
+        "--fps",
+        str(FPS),
+        "--num-frames",
+        str(NUM_FRAMES),
+        "--guidance-scale",
+        str(GUIDANCE_SCALE),
+        "--guidance-scale-2",
+        str(GUIDANCE_SCALE_2),
+        "--flow-shift",
+        str(FLOW_SHIFT),
+        "--num-inference-steps",
+        str(NUM_INFERENCE_STEPS),
+        "--seed",
+        str(SEED),
+        "--output",
+        str(output_path),
+        "--metadata-output",
+        str(metadata_path),
+    ]
+
+
+def _runner_env() -> dict[str, str]:
+    env = os.environ.copy()
+    pythonpath_parts = [
+        str(REPO_ROOT),
+        str(WORKSPACE_ROOT / "diffusers" / "src"),
+    ]
+    existing_pythonpath = env.get("PYTHONPATH")
+    if existing_pythonpath:
+        pythonpath_parts.append(existing_pythonpath)
+    env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
+    env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    return env
+
+
+def _resolve_image_source(configured: str | None) -> str:
+    configured = configured or RABBIT_IMAGE_URL
+    candidate = Path(configured)
+    if candidate.exists():
+        return str(candidate.resolve())
+    return configured
+
+
+def _is_remote_image_source(source: str) -> bool:
+    return source.startswith("http://") or source.startswith("https://")
+
+
+def _validate_image_source(source: str) -> None:
+    if _is_remote_image_source(source):
+        response = requests.get(source, timeout=60)
+        response.raise_for_status()
+        return
+
+    image_path = Path(source)
+    if not image_path.exists():
+        raise FileNotFoundError(f"Local image source does not exist: {image_path}")
+
+
+def _build_online_image_reference(source: str) -> str:
+    if _is_remote_image_source(source) or source.startswith("data:image"):
+        return source
+
+    image_path = Path(source)
+    mime_type = guess_type(image_path.name)[0] or "image/png"
+    encoded = b64encode(image_path.read_bytes()).decode("ascii")
+    return f"data:{mime_type};base64,{encoded}"
+
+
+def _probe_video(path: Path) -> dict[str, int | float]:
+    _probe_binary("ffprobe")
+    result = subprocess.run(
+        [
+            "ffprobe",
+            "-v",
+            "error",
+            "-select_streams",
+            "v:0",
+            "-count_frames",
+            "-show_entries",
+            "stream=width,height,avg_frame_rate,nb_read_frames",
+            "-of",
+            "json",
+            str(path),
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return _parse_video_metadata(json.loads(result.stdout))
+
+
+def _run_ffmpeg_similarity(filter_name: str, first: Path, second: Path) -> str:
+    _probe_binary("ffmpeg")
+    result = subprocess.run(
+        [
+            "ffmpeg",
+            "-hide_banner",
+            "-i",
+            str(first),
+            "-i",
+            str(second),
+            "-lavfi",
+            f"[0:v][1:v]{filter_name}",
+            "-f",
+            "null",
+            "-",
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return result.stderr
+
+
+def _online_timeout_seconds(configured: int | None) -> int:
+    return int(configured) if configured is not None else 1200
+
+
+def _artifact_dir(image_source: str) -> Path:
+    if _is_remote_image_source(image_source):
+        source_name = Path(urlparse(image_source).path).stem or "remote"
+    else:
+        source_name = Path(image_source).stem or "local"
+
+    safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
+    digest = sha1(image_source.encode("utf-8")).hexdigest()[:8]
+    return RESULT_ROOT / f"{safe_name}-{digest}"
+
+
+def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
+    artifact_dir = _artifact_dir(image_source)
+    artifact_dir.mkdir(parents=True, exist_ok=True)
+    return (
+        artifact_dir / "online.mp4",
+        artifact_dir / "offline.mp4",
+        artifact_dir / "offline_metadata.json",
+    )
+
+
+def _send_video_request_with_timeout(
+    openai_client,
+    request_config: dict[str, Any],
+    *,
+    timeout_seconds: int,
+) -> bytes:
+    form_data = request_config.get("form_data")
+    if not isinstance(form_data, dict):
+        raise ValueError("Video request_config must contain 'form_data'")
+
+    if not form_data.get("prompt"):
+        raise ValueError("Video request_config['form_data'] must contain 'prompt'")
+
+    normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
+
+    files: dict[str, tuple[str, BytesIO, str]] = {}
+    image_reference = request_config.get("image_reference")
+    if image_reference:
+        if image_reference.startswith("data:image"):
+            header, encoded = image_reference.split(",", 1)
+            content_type = header.split(";")[0].removeprefix("data:")
+            extension = content_type.split("/")[-1]
+            file_data = b64decode(encoded)
+            files["input_reference"] = (
+                f"reference.{extension}",
+                BytesIO(file_data),
+                content_type,
+            )
+        else:
+            normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
+
+    start_time = time.perf_counter()
+    create_url = openai_client._build_url("/v1/videos")
+    response = requests.post(
+        create_url,
+        data=normalized_form_data,
+        files=files,
+        headers={"Accept": "application/json"},
+        timeout=60,
+    )
+    response.raise_for_status()
+    job_data = response.json()
+    video_id = job_data["id"]
+    openai_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds)
+    video_content = openai_client._download_video_content(video_id)
+    print(f"online_video_e2e_latency_s={time.perf_counter() - start_time:.3f}")
+    return video_content
+
+
+def _generate_online_video(
+    *,
+    omni_server,
+    openai_client,
+    image_source: str,
+    online_timeout_seconds: int,
+) -> Path:
+    online_path, _, _ = _artifact_paths(image_source)
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "prompt": PROMPT,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "size": SIZE,
+            "fps": FPS,
+            "num_frames": NUM_FRAMES,
+            "guidance_scale": GUIDANCE_SCALE,
+            "guidance_scale_2": GUIDANCE_SCALE_2,
+            "flow_shift": FLOW_SHIFT,
+            "num_inference_steps": NUM_INFERENCE_STEPS,
+            "seed": SEED,
+        },
+        "image_reference": _build_online_image_reference(image_source),
+    }
+    online_video_bytes = _send_video_request_with_timeout(
+        openai_client,
+        request_config,
+        timeout_seconds=_online_timeout_seconds(online_timeout_seconds),
+    )
+    online_path.write_bytes(online_video_bytes)
+    return online_path
+
+
+def _generate_offline_video(*, image_source: str) -> tuple[Path, Path]:
+    _, offline_path, offline_metadata_path = _artifact_paths(image_source)
+    command = _build_diffusers_command(
+        runner_path=RUNNER_PATH,
+        image_source=image_source,
+        output_path=offline_path,
+        metadata_path=offline_metadata_path,
+    )
+    subprocess.run(
+        command,
+        cwd=REPO_ROOT,
+        env=_runner_env(),
+        check=True,
+        timeout=VIDEO_TIMEOUT_SECONDS,
+    )
+    return offline_path, offline_metadata_path
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_wan22_i2v_diffusers_offline_generates_video(
+    wan22_i2v_image_source: str | None,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("Wan2.2 I2V diffusers offline test requires CUDA.")
+
+    _probe_binary("ffprobe")
+    if not RUNNER_PATH.exists():
+        raise AssertionError(f"Offline diffusers runner does not exist: {RUNNER_PATH}")
+
+    image_source = _resolve_image_source(wan22_i2v_image_source)
+    _validate_image_source(image_source)
+    offline_path, offline_metadata_path = _generate_offline_video(image_source=image_source)
+    assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
+    assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}"
+    offline_metadata = _probe_video(offline_path)
+    assert offline_metadata["width"] == WIDTH
+    assert offline_metadata["height"] == HEIGHT
+    assert offline_metadata["fps"] == float(FPS)
+    assert offline_metadata["frame_count"] == NUM_FRAMES
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_wan22_i2v_online_serving_generates_video(
+    omni_server,
+    openai_client,
+    wan22_i2v_image_source: str | None,
+    wan22_i2v_online_timeout_seconds: int,
+) -> None:
+    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+        pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.")
+
+    _probe_binary("ffprobe")
+    image_source = _resolve_image_source(wan22_i2v_image_source)
+    _validate_image_source(image_source)
+    online_path = _generate_online_video(
+        omni_server=omni_server,
+        openai_client=openai_client,
+        image_source=image_source,
+        online_timeout_seconds=wan22_i2v_online_timeout_seconds,
+    )
+    assert online_path.exists(), f"Expected online video artifact at {online_path}"
+    online_metadata = _probe_video(online_path)
+    assert online_metadata["width"] == WIDTH
+    assert online_metadata["height"] == HEIGHT
+    assert online_metadata["fps"] == float(FPS)
+    assert online_metadata["frame_count"] == NUM_FRAMES
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_wan22_i2v_serving_matches_diffusers_video_similarity(
+    wan22_i2v_image_source: str | None,
+) -> None:
+    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+        pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.")
+
+    _probe_binary("ffmpeg")
+    _probe_binary("ffprobe")
+    if not RUNNER_PATH.exists():
+        raise AssertionError(f"Offline diffusers runner does not exist: {RUNNER_PATH}")
+
+    image_source = _resolve_image_source(wan22_i2v_image_source)
+    _validate_image_source(image_source)
+    online_path, offline_path, offline_metadata_path = _artifact_paths(image_source)
+
+    if not online_path.exists():
+        pytest.skip(f"Missing online artifact from prerequisite test: {online_path}")
+    if not offline_path.exists() or not offline_metadata_path.exists():
+        pytest.skip(f"Missing offline artifacts from prerequisite test: {offline_path}, {offline_metadata_path}")
+
+    assert online_path.exists(), f"Expected online video artifact at {online_path}"
+    assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
+    assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}"
+
+    online_metadata = _probe_video(online_path)
+    offline_metadata = _probe_video(offline_path)
+    assert online_metadata == offline_metadata, (
+        f"Video metadata mismatch:\n"
+        f"online={online_metadata}\n"
+        f"offline={offline_metadata}\n"
+        f"online_path={online_path}\n"
+        f"offline_path={offline_path}"
+    )
+    assert online_metadata["width"] == WIDTH
+    assert online_metadata["height"] == HEIGHT
+    assert online_metadata["fps"] == float(FPS)
+    assert online_metadata["frame_count"] == NUM_FRAMES
+
+    ssim_output = _run_ffmpeg_similarity("ssim", online_path, offline_path)
+    psnr_output = _run_ffmpeg_similarity("psnr", online_path, offline_path)
+    ssim_score = _parse_ssim_score(ssim_output)
+    psnr_score = _parse_psnr_score(psnr_output)
+
+    print("wan22_i2v_similarity metrics:")
+    print(
+        "  SSIM:"
+        f" value={ssim_score:.6f},"
+        f" threshold>={SSIM_THRESHOLD:.6f},"
+        " range=[0, 1],"
+        " higher_is_better=True,"
+        " interpretation=structural_similarity"
+    )
+    print(
+        "  PSNR:"
+        f" value={psnr_score:.6f} dB,"
+        f" threshold>={PSNR_THRESHOLD:.6f} dB,"
+        " range=[0, +inf),"
+        " higher_is_better=True,"
+        " interpretation=pixel_error_in_decibels"
+    )
+    print(f"online_video={online_path}")
+    print(f"offline_video={offline_path}")
+    print(f"offline_metadata={offline_metadata_path}")
+
+    assert ssim_score >= SSIM_THRESHOLD, (
+        f"SSIM below threshold: got {ssim_score:.6f}, expected >= {SSIM_THRESHOLD:.6f}. "
+        f"online={online_path} offline={offline_path}"
+    )
+    assert psnr_score >= PSNR_THRESHOLD, (
+        f"PSNR below threshold: got {psnr_score:.6f}, expected >= {PSNR_THRESHOLD:.6f}. "
+        f"online={online_path} offline={offline_path}"
+    )

tests/e2e/accuracy/wan22_i2v/wan22_i2v_video_similarity_common.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+RABBIT_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/rabbit.png"
+MODEL_NAME = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
+SIZE = "832x480"
+WIDTH = 832
+HEIGHT = 480
+FPS = 16
+NUM_FRAMES = 93
+GUIDANCE_SCALE = 3.5
+GUIDANCE_SCALE_2 = 3.5
+FLOW_SHIFT = 5.0
+NUM_INFERENCE_STEPS = 50
+SEED = 42
+BOUNDARY_RATIO = 0.875
+SSIM_THRESHOLD = 0.94
+PSNR_THRESHOLD = 28.0
+
+PROMPT = """一只棕色野兔的正面特写镜头，采用低角度仰拍视角，营造亲密而庄严的视觉冲击。兔子一双圆润漆黑的大眼睛直视镜头深处，眼神中交织着野生动物的警觉与一丝难以言喻的温柔好奇，仿佛在与观者建立跨越物种的静默对话。它毛色呈现层次丰富的棕褐渐变，从浅奶油色腹部过渡到深棕背部，每根毛发纹理清晰可辨，在侧光下泛着丝绸般的光泽。细长洁白的胡须共有三对，随呼吸节奏微微颤动，偶尔因捕捉气流信息而轻轻摇摆。
+两只标志性的长耳完全竖立，耳廓外侧覆盖短密棕毛，内侧则露出粉嫩的血管网络，薄如蝉翼的皮肤下血液流动隐约可见，耳朵以细微幅度不时转动，精准定位声源方向。背景是一片澄澈的蔚蓝天空，形态蓬松的白色积云以缓慢速度横向漂移，云影在兔子头顶交替变化，光线随之明暗流转。晴朗天气的明媚阳光从画面左上方45度角倾泻而下，在兔脸右侧形成柔和的伦勃朗式阴影，强化了面部立体感和皮毛质感。
+兔子湿润的黑鼻子持续进行每秒三至四次的快速抽动，这是它们感知化学信号的本能动作，粉色三瓣嘴随之轻启，露出正在反刍的洁白门齿，下颌以稳定节奏左右研磨。摄影采用大光圈浅景深，焦点牢牢锁定在兔子双眼连线所在的焦平面，背景天空和远景绿色植被虚化成圆润的彩色光斑，前景几根嫩绿草叶闯入画面边缘，以缓慢弧线随风摇曳，暗示着和煦的春日微风。
+整个场景弥漫着宁静致远的田园诗意，色彩温暖饱和，充满生命力。兔子在持续五秒的对视后，以典型 lagomorph 特征完成一次完整的瞬膜眨眼——第三眼睑从内侧横向滑过眼球，继而缓缓歪头向右十五度，这个行为在动物行为学中代表认知加工和好奇心表达，耳朵随之向同一方向倾斜，最终恢复正视姿态，胡须舒展，完成这段短暂而珍贵的自然纪录。"""
+
+NEGATIVE_PROMPT = """色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"""

tests/e2e/offline_inference/test_bagel_img2img.py

@@ -79,19 +79,17 @@ def _find_free_port() -> int:
     return port
 
 
-def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_steps: int = 15) -> list:
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
     """Configure sampling parameters for Bagel img2img generation.
 
     Args:
         omni: The Omni instance to get default params from.
-        max_tokens: Maximum tokens for the first stage.
         num_inference_steps: Number of inference steps for the diffusion stage.
 
     Returns:
         Configured sampling params list.
     """
     params_list = omni.default_sampling_params_list
-    params_list[0].max_tokens = max_tokens  # type: ignore
     if len(params_list) > 1:
         params_list[1].num_inference_steps = num_inference_steps  # type: ignore
         params_list[1].extra_args = {  # type: ignore

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -80,19 +80,17 @@ def _find_free_port() -> int:
     return port
 
 
-def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_steps: int = 15) -> list:
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
     """Configure sampling parameters for Bagel text2img generation.
 
     Args:
         omni: The Omni instance to get default params from.
-        max_tokens: Maximum tokens for the first stage.
         num_inference_steps: Number of inference steps for the diffusion stage.
 
     Returns:
         Configured sampling params list.
     """
     params_list = omni.default_sampling_params_list
-    params_list[0].max_tokens = max_tokens  # type: ignore
     if len(params_list) > 1:
         params_list[1].num_inference_steps = num_inference_steps  # type: ignore
         params_list[1].extra_args = {  # type: ignore

tests/e2e/offline_inference/test_quantization_fp8.py

@@ -120,7 +120,6 @@ def _generate_bagel_image(
         torch.cuda.reset_peak_memory_stats()
 
         params_list = omni.default_sampling_params_list
-        params_list[0].max_tokens = 1  # type: ignore
         if len(params_list) > 1:
             params_list[1].num_inference_steps = num_inference_steps  # type: ignore
             params_list[1].extra_args = {  # type: ignore

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -23,7 +23,7 @@ from tests.conftest import (
 )
 from tests.utils import hardware_test
 
-models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
+model = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
 
 AUDIO_KEY = ["test"]
 IMAGE_KEY = ["square", "quadrate"]
@@ -49,16 +49,32 @@ def get_chunk_config(default_path):
     return path
 
 
+def get_batch_token_config(default_path):
+    path = modify_stage_config(
+        default_path,
+        updates={
+            "stage_args": {1: {"engine_args.max_num_batched_tokens": 64}},
+        },
+    )
+    return path
+
+
 # CI stage config for 2*H100-80G GPUs
 default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
-stage_configs = [default_path, get_chunk_config(default_path)]
 
 if current_omni_platform.is_xpu():
-    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
+    default_path = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")
 
 # Create parameter combinations for model and stage config
 test_params = [
-    OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
+    pytest.param(OmniServerParams(model=model, stage_config_path=default_path), id="default"),
+    pytest.param(OmniServerParams(model=model, stage_config_path=get_chunk_config(default_path)), id="async_chunk"),
+]
+
+test_token_params = [
+    pytest.param(
+        OmniServerParams(model=model, stage_config_path=get_batch_token_config(default_path)), id="batch_token_64"
+    )
 ]
 
 
@@ -85,6 +101,8 @@ def get_prompt(prompt_type="text_only"):
         "text_video": "What is in this video? ",
         "text_image": "What is in this image? ",
         "text_audio": "What is in this audio? ",
+        "text_audio_video": "First, what is in this audio? Then, what is in this video? ",
+        "one_word": "What is the capital of UK? Answer in one word",
     }
     return prompts.get(prompt_type, prompts["text_only"])
 
@@ -121,7 +139,7 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.advanced_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
 def test_text_to_text_audio_001(omni_server, openai_client) -> None:
     """
     Input Modal: text
@@ -288,7 +306,7 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.advanced_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
 def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
     """
     Input Modal: text, audio
@@ -313,7 +331,7 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.advanced_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
 def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
     """
     Input Modal: text, image
@@ -339,7 +357,7 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.advanced_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
 def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
     """
     Input Modal: text, video
@@ -367,7 +385,7 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.advanced_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
 def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
     """
     Input Modal: text, audio, image, video
@@ -393,3 +411,106 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
         "key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY},
     }
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_in_video_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video).
+    Output Modal: text, audio
+    Input Setting: stream=False
+    Datasets: single request
+    """
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300, embed_audio=True)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        content_text=get_prompt("text_audio_video"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "use_audio_in_video": True,
+        "key_words": {"video": VIDEO_KEY, "audio": AUDIO_KEY + ["beep", "electronic"]},
+    }
+    openai_client.send_omni_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_in_video_002(omni_server, openai_client) -> None:
+    """
+    Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video).
+    Output Modal: text, audio
+    Input Setting: stream=True
+    Datasets: few requests
+    """
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300, embed_audio=True)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        content_text=get_prompt("text_audio_video"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "use_audio_in_video": True,
+        "key_words": {"video": VIDEO_KEY, "audio": AUDIO_KEY + ["beep", "electronic"]},
+    }
+
+    # Retry when assert_omni_response fails on key_words (see tests/conftest.py).
+    _keyword_assert_msg = "The output does not contain any of the keywords."
+    _max_retries = 3
+    for attempt in range(_max_retries):
+        try:
+            openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+            break
+        except AssertionError as e:
+            if _keyword_assert_msg not in str(e) or attempt == _max_retries - 1:
+                raise
+            print(f"Keyword assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_one_word_prompt_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text only (one-word answer constraint).
+    Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text=get_prompt("one_word"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "key_words": {"text": ["london"]},
+    }
+
+    # Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/conftest.py).
+    _similarity_assert_msg = "The audio content is not same as the text"
+    _max_retries = 3
+    for attempt in range(_max_retries):
+        try:
+            openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+            break
+        except AssertionError as e:
+            if _similarity_assert_msg not in str(e) or attempt == _max_retries - 1:
+                raise
+            print(f"Similarity assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")

tests/e2e/online_serving/test_qwen_image_expansion.py

@@ -28,6 +28,11 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
 
 def _get_diffusion_feature_cases(model: str):
     return [
+        pytest.param(
+            OmniServerParams(model=model, server_args=["--step-execution"]),
+            id="step_execution",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
         pytest.param(
             OmniServerParams(model=model, server_args=["--cache-backend", "tea_cache"]),
             id="cache_tea_cache",

tests/e2e/online_serving/test_qwen_image_layered_expansion.py

@@ -7,7 +7,7 @@ and are supported by Qwen-Image-Layered model.
 Kept cases (maximum feature coverage, no cpu-offload):
   sp_001              : cache_dit + Ulysses-SP 2          (2×H100)
   cfg_parallel_002    : cache_dit + CFG-Parallel 2        (2×H100)
-  layers_guard_001    : layerwise offload + layers=2      (1×H100, issue #1969 guard)
+  layers_guard_001    : layerwise offload + layers=3      (1×H100, issue #1969 guard)
 
 Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload.
 """
@@ -102,8 +102,8 @@ LAYERS_GUARD_CASES = [
             model=MODEL,
             server_args=["--enable-layerwise-offload"],
         ),
-        2,
-        id="layers_guard_001_layers2",
+        3,
+        id="layers_guard_001_layers3",
         marks=SINGLE_CARD_FEATURE_MARKS,
     ),
 ]

tests/e2e/stage_configs/qwen3_omni_ci.yaml

@@ -33,7 +33,7 @@ stage_args:
     temperature: 0.4
     top_p: 0.9
     top_k: 1
-    max_tokens: 100
+    max_tokens: 150
     seed: 42
     ignore_eos: False
     detokenize: True
@@ -74,7 +74,7 @@ stage_args:
     devices: "1"
   engine_args:
     model_stage: code2wav
-    max_num_seqs: 1
+    max_num_seqs: 5
     model_arch: Qwen3OmniMoeForConditionalGeneration
     worker_type: generation
     scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler

tests/entrypoints/openai_api/test_video_server.py

@@ -735,3 +735,156 @@ def test_extra_params_merged_with_existing_extra_args(test_client, mocker: Mocke
     assert captured.extra_args["flow_shift"] == 0.5
     assert captured.extra_args["use_zero_init"] is True
     assert captured.extra_args["zero_steps"] == 2
+
+
+# ---------------------------------------------------------------------------
+# Sync endpoint tests (POST /v1/videos/sync)
+# ---------------------------------------------------------------------------
+
+
+def _mock_encode_video_bytes(mocker: MockerFixture, return_value: bytes = b"fake-video-bytes"):
+    """Mock the raw-bytes encoder used by the sync video path."""
+    return mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=return_value,
+    )
+
+
+def test_sync_t2v_returns_video_bytes(test_client, mocker: MockerFixture):
+    """Sync endpoint should block until generation finishes and return raw
+    video bytes with metadata headers."""
+    _mock_encode_video_bytes(mocker, b"fake-video-bytes")
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={
+            "prompt": "A cat running across the street.",
+            "size": "640x360",
+            "seconds": "2",
+            "fps": "12",
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "video/mp4"
+    assert response.content == b"fake-video-bytes"
+    assert response.headers["x-request-id"].startswith("video_sync-")
+    assert response.headers["x-model"] == "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+    assert float(response.headers["x-inference-time-s"]) >= 0
+
+
+def test_sync_i2v_returns_video_bytes(test_client, mocker: MockerFixture):
+    """Sync I2V endpoint should accept an uploaded reference image and return
+    raw video bytes."""
+    image_bytes = _make_test_image_bytes((48, 32))
+    _mock_encode_video_bytes(mocker, b"i2v-video-data")
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={"prompt": "A bear playing with yarn."},
+        files={"input_reference": ("input.png", image_bytes, "image/png")},
+    )
+
+    assert response.status_code == 200
+    assert response.content == b"i2v-video-data"
+    assert response.headers["content-type"] == "video/mp4"
+
+
+def test_sync_i2v_with_image_reference(test_client, mocker: MockerFixture):
+    """Sync I2V endpoint should accept a JSON image_reference field."""
+    _mock_encode_video_bytes(mocker, b"ref-video")
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={
+            "prompt": "A fox running through snow.",
+            "image_reference": json.dumps({"image_url": _make_test_image_data_url((40, 24))}),
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.content == b"ref-video"
+
+
+def test_sync_missing_handler_returns_503():
+    app = FastAPI()
+    app.include_router(router)
+    app.state.openai_serving_video = None
+    client = TestClient(app)
+
+    response = client.post(
+        "/v1/videos/sync",
+        data={"prompt": "no handler"},
+    )
+    assert response.status_code == 503
+    assert "not initialized" in response.json()["detail"].lower()
+
+
+def test_sync_missing_prompt_returns_422(test_client):
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={"size": "320x240"},
+    )
+    assert response.status_code == 422
+
+
+def test_sync_rejects_both_references(test_client):
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={
+            "prompt": "bad refs",
+            "image_reference": '{"image_url": "https://example.com/cat.png"}',
+        },
+        files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
+    )
+    assert response.status_code == 400
+    assert "either input_reference or image_reference" in response.json()["detail"].lower()
+
+
+def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture):
+    """If the underlying generation raises, the sync endpoint should return 500."""
+    mocker.patch.object(
+        OmniOpenAIServingVideo,
+        "generate_video_bytes",
+        side_effect=RuntimeError("GPU exploded"),
+    )
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={"prompt": "will fail"},
+    )
+    assert response.status_code == 500
+    assert "GPU exploded" in response.json()["detail"]
+
+
+def test_sync_does_not_create_store_entry(test_client, mocker: MockerFixture):
+    """The sync endpoint should NOT leave any record in VIDEO_STORE — it is
+    stateless by design."""
+    _mock_encode_video_bytes(mocker)
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={"prompt": "stateless test"},
+    )
+    assert response.status_code == 200
+    loop = asyncio.new_event_loop()
+    try:
+        stored = loop.run_until_complete(api_server.VIDEO_STORE.list_values())
+    finally:
+        loop.close()
+    assert len(stored) == 0
+
+
+def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
+    """Sampling parameters should propagate to the engine through the sync path."""
+    _mock_encode_video_bytes(mocker)
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={
+            "prompt": "param pass",
+            "num_inference_steps": "30",
+            "guidance_scale": "6.5",
+            "seed": "42",
+        },
+    )
+    assert response.status_code == 200
+    engine = test_client.app.state.openai_serving_video._engine_client
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.num_inference_steps == 30
+    assert captured.guidance_scale == 6.5
+    assert captured.seed == 42

tests/entrypoints/test_async_omni_abort.py

@@ -0,0 +1,85 @@
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_generate_accepts_request_after_repeated_cancellations():
+    async def run_test():
+        submitted_request_ids = []
+        aborted_request_batches = []
+
+        async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id):
+            del prompt, sampling_params_list, final_stage_id
+            submitted_request_ids.append(request_id)
+
+        async def fake_abort_async(request_ids):
+            aborted_request_batches.append(list(request_ids))
+
+        async def fake_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts):
+            del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts
+            if request_id.startswith("cancel-"):
+                await asyncio.Future()
+                return
+            yield SimpleNamespace(
+                stage_id=0,
+                request_output=SimpleNamespace(outputs=[]),
+                finished=True,
+            )
+
+        async def collect_outputs(request_id):
+            outputs = []
+            async for output in AsyncOmni.generate(
+                omni,
+                prompt={"prompt": "prompt"},
+                request_id=request_id,
+                sampling_params_list=[SimpleNamespace()],
+                output_modalities=["image"],
+            ):
+                outputs.append(output)
+            return outputs
+
+        omni = object.__new__(AsyncOmni)
+        omni._pause_cond = asyncio.Condition()
+        omni._paused = False
+        omni.engine = SimpleNamespace(
+            num_stages=1,
+            add_request_async=fake_add_request_async,
+            abort_async=fake_abort_async,
+        )
+        omni.log_stats = False
+        omni.request_states = {}
+        omni._final_output_handler = lambda: None
+        omni.resolve_sampling_params_list = lambda params: params
+        omni._compute_final_stage_id = lambda output_modalities: 0
+        omni._process_orchestrator_results = fake_process_results
+        omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)
+
+        assert len(await collect_outputs("baseline")) == 1
+
+        for idx in range(3):
+            task = asyncio.create_task(collect_outputs(f"cancel-{idx}"))
+            await asyncio.sleep(0)
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+
+        assert len(await collect_outputs("after-cancel")) == 1
+        assert submitted_request_ids == [
+            "baseline",
+            "cancel-0",
+            "cancel-1",
+            "cancel-2",
+            "after-cancel",
+        ]
+        assert aborted_request_batches == [
+            ["cancel-0"],
+            ["cancel-1"],
+            ["cancel-2"],
+        ]
+
+    asyncio.run(run_test())

tests/entrypoints/test_async_omni_diffusion.py

@@ -1,9 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import asyncio
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from types import SimpleNamespace
+from unittest.mock import Mock
+
 import pytest
 
+import vllm_omni.diffusion.stage_diffusion_client as stage_diffusion_client_module
+from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
 from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -13,3 +23,91 @@ def test_get_diffusion_od_config_returns_direct_config():
     diffusion.od_config = object()
 
     assert diffusion.get_diffusion_od_config() is diffusion.od_config
+
+
+def test_async_omni_diffusion_generate_aborts_engine_on_cancel():
+    async def run_test():
+        started = threading.Event()
+        release = threading.Event()
+        abort = Mock()
+
+        def step(request):
+            del request
+            started.set()
+            release.wait(timeout=5)
+            return [SimpleNamespace(request_id="req-1")]
+
+        diffusion = object.__new__(AsyncOmniDiffusion)
+        diffusion.engine = SimpleNamespace(step=step, abort=abort)
+        diffusion._executor = ThreadPoolExecutor(max_workers=1)
+
+        task = asyncio.create_task(
+            diffusion.generate(
+                prompt="hello",
+                sampling_params=OmniDiffusionSamplingParams(),
+                request_id="req-1",
+            )
+        )
+        try:
+            assert await asyncio.to_thread(started.wait, 1)
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+        finally:
+            release.set()
+            diffusion._executor.shutdown(wait=True)
+
+        abort.assert_called_once_with("req-1")
+
+    asyncio.run(run_test())
+
+
+def test_stage_diffusion_client_abort_requests_forwards_to_engine():
+    async def run_test():
+        aborted_request_ids: list[list[str]] = []
+
+        async def abort(request_ids):
+            aborted_request_ids.append(request_ids)
+
+        client = object.__new__(StageDiffusionClient)
+        client._engine = SimpleNamespace(abort=abort)
+        client._tasks = {}
+
+        task = asyncio.create_task(asyncio.sleep(60))
+        client._tasks["req-1"] = task
+
+        await client.abort_requests_async(["req-1", "req-2"])
+
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        assert client._tasks == {}
+        assert aborted_request_ids == [["req-1", "req-2"]]
+
+    asyncio.run(run_test())
+
+
+def test_stage_diffusion_client_run_treats_abort_as_normal_path(monkeypatch):
+    async def run_test():
+        async def generate(prompt, sampling_params, request_id):
+            del prompt, sampling_params
+            raise DiffusionRequestAbortedError(f"Request {request_id} aborted.")
+
+        info = Mock()
+        exception = Mock()
+        monkeypatch.setattr(stage_diffusion_client_module.logger, "info", info)
+        monkeypatch.setattr(stage_diffusion_client_module.logger, "exception", exception)
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 3
+        client._engine = SimpleNamespace(generate=generate)
+        client._output_queue = asyncio.Queue()
+        client._tasks = {"req-1": object()}
+
+        await client._run("req-1", "prompt", OmniDiffusionSamplingParams())
+
+        assert client._output_queue.empty()
+        assert client._tasks == {}
+        info.assert_called_once()
+        exception.assert_not_called()
+
+    asyncio.run(run_test())

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -16,6 +16,8 @@ from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
     talker2code2wav_async_chunk,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 _FRAME = [1, 2, 3, 4]
 _Q = len(_FRAME)
 
@@ -90,26 +92,55 @@ def test_flush_on_finish():
         is_finished=True,
     )
     assert p is not None
-    assert p["finished"].item() is True
+    assert p["finished"] is True
     assert len(p["code_predictor_codes"]) == _Q * 24
 
 
 _CASES = [
+    # ── IC boundary rule ──────────────────────────────────────────────
+    # IC phase: length <= chunk_size  (uses <=, consistent with fish_speech)
+    # IC emits fill the entire first chunk_size worth of frames, so the
+    # normal phase always starts at a clean chunk boundary.
+    # initial_coverage = (chunk_size // initial_chunk_size) * initial_chunk_size
+    #
     # Dynamic IC=16, cs=25, initial_coverage=16
-    ((25, 25, 0), 24, False, None),  # IC phase: 24%16!=0 -> hold
-    ((25, 25, 0), 25, False, None),  # transition: adjusted=9, hold (no replay)
-    ((25, 25, 0), 41, False, (16, 41)),  # first normal emit, lc=16
+    # IC does NOT evenly divide cs, so initial_coverage < cs.
+    # IC emits at 16; frames 17-25 remain in IC phase but 25%16!=0 -> hold.
+    # Normal phase: adjusted = length - 16, emit when adjusted % 25 == 0.
+    ((25, 25, 0), 24, False, None),  # IC: 24<=25, 24%16!=0 -> hold
+    ((25, 25, 0), 25, False, None),  # IC: 25<=25, 25%16!=0 -> hold
+    ((25, 25, 0), 41, False, (16, 41)),  # normal: adjusted=25, 25%25==0 -> emit, lc=16
+    #
     # Per-request IC=10, cs=25, initial_coverage=20
-    ((25, 25, 10), 9, False, None),  # IC: hold
-    ((25, 25, 10), 10, False, (0, 10)),  # IC: emit at boundary
-    ((25, 25, 10), 25, False, None),  # transition: hold (no replay)
-    ((25, 25, 10), 45, False, (20, 45)),  # first normal emit, lc=20
+    # IC does NOT evenly divide cs; IC emits at 10, 20.
+    # Frames 21-25 are still IC phase but 21..25 % 10 != 0 -> hold.
+    ((25, 25, 10), 9, False, None),  # IC: 9%10!=0 -> hold
+    ((25, 25, 10), 10, False, (0, 10)),  # IC: 10%10==0 -> emit, lc=0
+    ((25, 25, 10), 25, False, None),  # IC: 25<=25, 25%10!=0 -> hold
+    ((25, 25, 10), 45, False, (20, 45)),  # normal: adjusted=25, 25%25==0 -> emit, lc=20
     ((25, 25, 10), 5, True, (0, 5)),  # finished flushes IC tail
     ((25, 25, 10), 33, True, (20, 33)),  # finished flushes normal tail
-    # IC=8, cs=16: IC divides chunk_size evenly (edge case)
-    ((16, 25, 8), 8, False, (0, 8)),  # IC: emit
-    ((16, 25, 8), 16, False, None),  # transition: hold (no replay)
-    ((16, 25, 8), 24, False, (8, 24)),  # first normal emit, lc=8
+    #
+    # IC=8, cs=16: IC evenly divides chunk_size (edge case)
+    # initial_coverage = (16//8)*8 = 16 == chunk_size.
+    # IC fills the entire first chunk: emits at 8 and 16.
+    # Normal phase starts at frame 17; first normal emit at 16+16=32.
+    ((16, 25, 8), 8, False, (0, 8)),  # IC: 8%8==0 -> emit, lc=0
+    ((16, 25, 8), 16, False, (8, 16)),  # IC: 16<=16, 16%8==0 -> emit, lc=8
+    ((16, 25, 8), 24, False, None),  # normal: adjusted=8, 8%16!=0 -> hold
+    ((16, 25, 8), 32, False, (16, 32)),  # normal: adjusted=16, 16%16==0 -> first emit, lc=16
+    #
+    # IC=5, cs=25: IC evenly divides chunk_size
+    # initial_coverage = (25//5)*5 = 25 == chunk_size.
+    # IC fills the entire first chunk: emits at 5, 10, 15, 20, 25.
+    # Normal phase starts at frame 26; first normal emit at 25+25=50.
+    # Emit intervals: 5,5,5,5,5,25,25,... — smooth transition, no gap.
+    ((25, 25, 5), 5, False, (0, 5)),  # IC: 5%5==0 -> emit, lc=0
+    ((25, 25, 5), 12, False, None),  # IC: 12%5!=0 -> hold
+    ((25, 25, 5), 25, False, (20, 25)),  # IC: 25<=25, 25%5==0 -> emit, lc=20
+    ((25, 25, 5), 30, False, None),  # normal: adjusted=5, 5%25!=0 -> hold
+    ((25, 25, 5), 50, False, (25, 50)),  # normal: adjusted=25, 25%25==0 -> first emit, lc=25
+    #
     # Per-request override: IC=15 at n_frames=10 -> 10%15!=0 -> hold
     ((25, 25, 15), 10, False, None),
 ]
@@ -158,24 +189,29 @@ def test_dynamic_ic_adapts_to_load():
 
 
 def test_ic_load_change_mid_request():
-    """IC stateless: load spike mid-request shifts initial_coverage."""
+    """IC is cached per request; a load spike only affects new requests."""
     tm = _tm(chunk_frames=25, left_context=25, max_num_seqs=8)
 
-    # Low load -> IC=2 -> emit at frame 2
+    # Low load -> IC=2 (cached for "r"), emit at frame 2
     p1 = _call(tm, "r", n_frames=2)
     assert p1 is not None
 
-    # Spike load: 6 others -> IC=16 -> initial_coverage=16
+    # Spike load: 6 others running
     for i in range(6):
         tm.code_prompt_token_ids[f"other-{i}"] = [[0]] * 10
 
-    # adjusted=25-16=9, 9%25!=0 -> hold
+    # IC for "r" is still cached as 2.
+    # initial_coverage = ((25-1)//2)*2 = 24, first normal emit at 24+25=49
     assert _call(tm, "r", n_frames=25) is None
-
-    # First normal emit at 16+25=41
-    p3 = _call(tm, "r", n_frames=41)
+    assert _call(tm, "r", n_frames=27) is None
+    p3 = _call(tm, "r", n_frames=49)
     assert p3 is not None
-    assert p3["left_context_size"] == 16
+
+    # A *new* request under high load gets IC=16 (not IC=2).
+    # Frame 2 would emit under IC=2 but must hold under IC=16.
+    assert _call(tm, "new_req", n_frames=2) is None
+    p4 = _call(tm, "new_req", n_frames=16)
+    assert p4 is not None
 
 
 @pytest.mark.parametrize(
@@ -295,8 +331,13 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
         ],
         dtype=torch.long,
     )
-    output = SimpleNamespace(multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code})
-    stage = SimpleNamespace(engine_outputs=[SimpleNamespace(outputs=[output])])
+    output = SimpleNamespace(
+        multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code},
+        token_ids=list(range(3)),
+    )
+    stage = SimpleNamespace(
+        engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
+    )
 
     prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])
 
@@ -335,8 +376,13 @@ def test_non_async_processor_filters_out_of_range_codec_values():
         ],
         dtype=torch.long,
     )
-    output = SimpleNamespace(multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code})
-    stage = SimpleNamespace(engine_outputs=[SimpleNamespace(outputs=[output])])
+    output = SimpleNamespace(
+        multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code},
+        token_ids=list(range(4)),
+    )
+    stage = SimpleNamespace(
+        engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
+    )
 
     prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])
 

tests/utils.py

@@ -430,13 +430,41 @@ def xpu_marks(*, res: str, num_cards: int):
         return marks + [test_distributed]
 
 
+def musa_marks(*, res: str, num_cards: int):
+    """
+    Get a collection of pytest marks to apply for `@musa_test`.
+
+    Args:
+        res: Resource type, e.g., "S5000".
+        num_cards: Number of GPU cards required.
+
+    Returns:
+        List of pytest marks to apply.
+    """
+    test_platform_detail = pytest.mark.musa
+
+    if res == "S5000":
+        test_resource = pytest.mark.S5000
+    else:
+        raise ValueError(f"Invalid MUSA resource type: {res}. Supported: S5000")
+
+    marks = [test_resource, test_platform_detail]
+
+    if num_cards == 1:
+        return marks
+    else:
+        test_distributed = pytest.mark.distributed_musa(num_cards=num_cards)
+        # TODO: add MUSA support for `skipif_musa` marker
+        return marks + [test_distributed]
+
+
 def gpu_marks(*, res: str, num_cards: int):
     """
     Get a collection of pytest marks to apply for `@gpu_test`.
     Platform is automatically determined based on resource type.
 
     Args:
-        res: Resource type, e.g., "L4", "H100" for CUDA, or "MI325" for ROCm, or "B60" for XPU.
+        res: Resource type, e.g., "L4", "H100" for CUDA, or "MI325" for ROCm, or "B60" for XPU, or "S5000" for MUSA.
         num_cards: Number of GPU cards required.
 
     Returns:
@@ -449,7 +477,9 @@ def gpu_marks(*, res: str, num_cards: int):
         return [test_platform] + rocm_marks(res=res, num_cards=num_cards)
     if res == "B60":
         return [test_platform] + xpu_marks(res=res, num_cards=num_cards)
-    raise ValueError(f"Invalid resource type: {res}. Supported: L4, H100, MI325")
+    if res == "S5000":
+        return [test_platform] + musa_marks(res=res, num_cards=num_cards)
+    raise ValueError(f"Invalid resource type: {res}. Supported: L4, H100, MI325, B60, S5000")
 
 
 def npu_marks(*, res: str, num_cards: int):
@@ -476,13 +506,13 @@ def npu_marks(*, res: str, num_cards: int):
 def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
     """
     Get a collection of pytest marks to apply for `@hardware_test`,
-    including CUDA, ROCm, XPU, and NPU,
+    including CUDA, ROCm, XPU, NPU, and MUSA,
     based on the specified platforms and resources.
     """
     # Validate platforms
     # Don't validate platform details in this decorator
     for platform, _ in res.items():
-        if platform not in ("cuda", "rocm", "xpu", "npu"):
+        if platform not in ("cuda", "rocm", "xpu", "npu", "musa"):
             raise ValueError(f"Unsupported platform: {platform}")
 
     # Normalize num_cards
@@ -505,6 +535,8 @@ def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
         cards = num_cards_dict[platform]
         if platform == "cuda" or platform == "rocm" or platform == "xpu":
             marks = gpu_marks(res=resource, num_cards=cards)
+        elif platform == "musa":
+            marks = musa_marks(res=resource, num_cards=cards)
         elif platform == "npu":
             marks = npu_marks(res=resource, num_cards=cards)
         else:
@@ -522,15 +554,17 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
         res: Mapping from platform to resource type. Supported platforms/resources:
             - cuda: L4, H100
             - rocm: MI325
+            - xpu: B60
             - npu: A2, A3
+            - musa: S5000
         num_cards: Number of cards required. Can be:
             - int: same card count for all platforms (default: 1)
             - dict: per-platform card count, e.g., {"cuda": 2, "rocm": 2}
 
     Example:
         @hardware_test(
-            res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
-            num_cards={"cuda": 2, "rocm": 2, "npu": 2},
+            res={"cuda": "L4", "rocm": "MI325", "npu": "A2", "musa": "S5000"},
+            num_cards={"cuda": 2, "rocm": 2, "npu": 2, "musa": 2},
         )
         def test_multi_platform():
             ...

tools/nightly/generate_nightly_perf_excel.py

@@ -78,16 +78,22 @@ BENCHMARK_COLUMNS: tuple[str, ...] = (
     "output_throughput",
     "total_token_throughput",
     "mean_ttft_ms",
+    "median_ttft_ms",
     "p99_ttft_ms",
     "mean_tpot_ms",
+    "median_tpot_ms",
     "p99_tpot_ms",
     "mean_itl_ms",
+    "median_itl_ms",
     "p99_itl_ms",
     "mean_e2el_ms",
+    "median_e2el_ms",
     "p99_e2el_ms",
     "mean_audio_rtf",
+    "median_audio_rtf",
     "p99_audio_rtf",
     "mean_audio_duration_s",
+    "median_audio_duration_s",
     "p99_audio_duration_s",
 )
 # Columns that get float coercion and number format in Excel. Excludes request_rate ("inf" str)
@@ -143,16 +149,22 @@ def _load_summary_columns(script_dir: str) -> list[str]:
         "output_throughput",
         "total_token_throughput",
         "mean_ttft_ms",
+        "median_ttft_ms",
         "p99_ttft_ms",
         "mean_tpot_ms",
+        "median_tpot_ms",
         "p99_tpot_ms",
         "mean_itl_ms",
+        "median_itl_ms",
         "p99_itl_ms",
         "mean_e2el_ms",
+        "median_e2el_ms",
         "p99_e2el_ms",
         "mean_audio_rtf",
+        "median_audio_rtf",
         "p99_audio_rtf",
         "mean_audio_duration_s",
+        "median_audio_duration_s",
         "p99_audio_duration_s",
         "commit_sha",
         "build_id",
@@ -447,14 +459,29 @@ def _apply_build_metadata_to_latest_only(
     build_id: str | None,
     build_url: str | None,
 ) -> None:
-    """Set commit_sha, build_id, build_url only on rows with the latest date.
-    Other rows get None so that build info is not duplicated for older benchmark data.
+    """Set commit_sha, build_id, build_url on rows from the latest calendar day.
+
+    Dates are expected like YYYYMMDD-HHMMSS (filename / benchmark convention). All rows
+    whose date starts with the same YYYYMMDD as the lexicographic max date receive
+    build metadata; older calendar days get None.
+    When max date is shorter than 8 chars, falls back to exact match.
     """
     if not records:
         return
     max_date = max((r.get("date") or "") for r in records)
+    use_day_prefix = len(max_date) >= 8
+    day_prefix = max_date[:8] if use_day_prefix else ""
+
     for r in records:
-        if (r.get("date") or "") == max_date:
+        d = r.get("date") or ""
+        if use_day_prefix and d.startswith(day_prefix):
+            in_latest_day = True
+        elif not use_day_prefix and d == max_date:
+            in_latest_day = True
+        else:
+            in_latest_day = False
+
+        if in_latest_day:
             r["commit_sha"] = commit_sha
             r["build_id"] = build_id
             r["build_url"] = build_url

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -95,8 +95,19 @@ class OmniARScheduler(VLLMScheduler):
             return False
 
         criteria_type = self.kv_transfer_criteria.get("type")
+        if (
+            self.kv_transfer_criteria.get("stop_after_transfer", True)
+            and request.request_id in self.transfer_triggered_requests
+        ):
+            # For split pipelines that only need the transferred KV
+            # snapshot, stop AR decode once KV extraction has completed.
+            # This frees stage-0 resources without requiring an
+            # orchestrator-side abort.
+            if request.request_id not in self.active_kv_transfers:
+                request.status = RequestStatus.FINISHED_STOPPED
+                return True
+            return False
 
-        # Universal duplicate check for once semantics
         if request.request_id in self.transfer_triggered_requests:
             return False
 
@@ -456,6 +467,23 @@ class OmniARScheduler(VLLMScheduler):
             kv_extracted_ids = getattr(model_runner_output, "kv_extracted_req_ids", None)
             if kv_extracted_ids:
                 for req_id in kv_extracted_ids:
+                    # Emit a kv_ready signal so the orchestrator can forward
+                    # the request to the DiT stage immediately after KV
+                    # extraction, without waiting for AR decode to finish.
+                    req = self.requests.get(req_id)
+                    if req is not None and not req.is_finished():
+                        eco = engine_core_outputs.get(req.client_index)
+                        if eco is None:
+                            eco = EngineCoreOutputs()
+                            engine_core_outputs[req.client_index] = eco
+                        eco.outputs.append(
+                            EngineCoreOutput(
+                                request_id=req_id,
+                                new_token_ids=[],
+                                kv_transfer_params={"kv_ready": True},
+                            )
+                        )
+
                     # Mark transfer as finished
                     if req_id in self.active_kv_transfers:
                         self.active_kv_transfers.remove(req_id)

vllm_omni/diffusion/attention/backends/abstract.py

@@ -99,6 +99,8 @@ class AttentionImpl(ABC, Generic[T]):
             return self.forward_npu(query, key, value, attn_metadata)
         elif current_omni_platform.is_xpu():
             return self.forward_xpu(query, key, value, attn_metadata)
+        elif current_omni_platform.is_musa():
+            return self.forward_musa(query, key, value, attn_metadata)
         else:
             raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}")
 
@@ -138,3 +140,13 @@ class AttentionImpl(ABC, Generic[T]):
     ) -> torch.Tensor:
         # By default, HIP ops are compatible with CUDA ops.
         return self.forward_cuda(query, key, value, attn_metadata)
+
+    def forward_musa(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None = None,
+    ) -> torch.Tensor:
+        # By default, MUSA ops are compatible with CUDA ops.
+        return self.forward_cuda(query, key, value, attn_metadata)

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

@@ -35,6 +35,9 @@ elif current_omni_platform.is_xpu():
         from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func  # noqa: F401
     except (ImportError, ModuleNotFoundError):
         pass
+elif current_omni_platform.is_musa():
+    # XXX (MUSA): Add MUSA-specific Flash Attention when available
+    pass
 else:
     # CUDA: try FA3 -> FA2 fallback chain
     # Try FA3 from fa3-fwd PyPI package

vllm_omni/diffusion/data.py

@@ -492,6 +492,9 @@ class OmniDiffusionConfig:
     # Step mode settings
     step_execution: bool = False
 
+    # Maximum number of sequences to generate in a batch
+    max_num_seqs: int = 1
+
     @property
     def is_moe(self) -> bool:
         num_experts = self.tf_model_config.get("num_experts", None)
@@ -658,6 +661,8 @@ class DiffusionOutput:
     trajectory_latents: torch.Tensor | None = None
     trajectory_decoded: list[torch.Tensor] | None = None
     error: str | None = None
+    aborted: bool = False
+    abort_message: str | None = None
 
     post_process_func: Callable[..., Any] | None = None
 
@@ -675,6 +680,10 @@ class DiffusionOutput:
     peak_memory_mb: float = 0.0
 
 
+class DiffusionRequestAbortedError(RuntimeError):
+    """Raised when a diffusion request ends via user-visible abort."""
+
+
 class AttentionBackendEnum(enum.Enum):
     FA = enum.auto()
     SLIDING_TILE_ATTN = enum.auto()

vllm_omni/diffusion/diffusion_engine.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from __future__ import annotations
+
+import queue
 import threading
 import time
 from collections.abc import Iterable
@@ -11,7 +14,11 @@ import PIL.Image
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import (
+    DiffusionOutput,
+    DiffusionRequestAbortedError,
+    OmniDiffusionConfig,
+)
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.registry import (
     DiffusionModelRegistry,
@@ -19,7 +26,9 @@ from vllm_omni.diffusion.registry import (
     get_diffusion_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface
+from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler
+from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
+from vllm_omni.diffusion.worker.utils import RunnerOutput
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -72,9 +81,14 @@ class DiffusionEngine:
 
         executor_class = DiffusionExecutor.get_class(od_config)
         self.executor = executor_class(od_config)
-        self.scheduler: SchedulerInterface = scheduler or RequestScheduler()
+        self.step_execution = bool(getattr(od_config, "step_execution", False))
+        self.scheduler: SchedulerInterface = scheduler or (
+            StepScheduler() if self.step_execution else RequestScheduler()
+        )
         self.scheduler.initialize(od_config)
-        self._rpc_lock = threading.Lock()
+        self._rpc_lock = threading.RLock()
+        self.abort_queue: queue.Queue[str] = queue.Queue()
+        self.execute_fn = self.executor.execute_step if self.step_execution else self.executor.execute_request
 
         try:
             self._dummy_run()
@@ -98,6 +112,8 @@ class DiffusionEngine:
         output = self.add_req_and_wait_for_response(request)
         exec_total_time = time.perf_counter() - exec_start_time
 
+        if output.aborted:
+            raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.")
         if output.error:
             raise Exception(f"{output.error}")
         logger.info("Generation completed successfully.")
@@ -264,7 +280,7 @@ class DiffusionEngine:
     def make_engine(
         config: OmniDiffusionConfig,
         scheduler: SchedulerInterface | None = None,
-    ) -> "DiffusionEngine":
+    ) -> DiffusionEngine:
         """Factory method to create a DiffusionEngine instance.
 
         Args:
@@ -281,8 +297,11 @@ class DiffusionEngine:
 
             # keep scheduling and executing until the target request is finished
             while True:
+                self._process_aborts_queue()
                 sched_output = self.scheduler.schedule()
                 if sched_output.is_empty:
+                    if target_sched_req_id in sched_output.finished_req_ids:
+                        return self._finalize_finished_request(target_sched_req_id)
                     if not self.scheduler.has_requests():
                         raise RuntimeError("Diffusion scheduler has no runnable requests.")
                     continue
@@ -292,21 +311,26 @@ class DiffusionEngine:
                 # vllm_omni/diffusion/sched/base_scheduler.py), so we directly
                 # take the single scheduled request here.
                 sched_req_id = sched_output.scheduled_req_ids[0]
-                req = sched_output.scheduled_new_reqs[0].req
                 try:
-                    output = self.executor.add_req(req)
+                    runner_output = self.execute_fn(sched_output)
                 except Exception as exc:
-                    logger.error(
-                        "Execution failed for diffusion request %s",
-                        sched_req_id,
-                        exc_info=True,
+                    logger.error("Execution failed for diffusion request %s", sched_req_id, exc_info=True)
+                    runner_output = RunnerOutput(
+                        req_id=sched_req_id,
+                        step_index=None,
+                        finished=True,
+                        result=DiffusionOutput(error=str(exc)),
                     )
-                    output = DiffusionOutput(error=str(exc))
 
-                finished_req_ids = self.scheduler.update_from_output(sched_output, output)
+                self._process_aborts_queue()
+
+                finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
                 if target_sched_req_id in finished_req_ids:
-                    self.scheduler.pop_request_state(target_sched_req_id)
-                    return output
+                    return self._finalize_finished_request(
+                        target_sched_req_id,
+                        runner_output=runner_output,
+                        missing_result_error="Diffusion execution finished without a final output.",
+                    )
 
     def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
         """Start or stop torch profiling on all diffusion workers.
@@ -437,6 +461,55 @@ class DiffusionEngine:
             self.executor.shutdown()
 
     def abort(self, request_id: str | Iterable[str]) -> None:
-        # TODO implement it
-        logger.warning("DiffusionEngine abort is not implemented yet")
-        pass
+        request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
+        for req_id in request_ids:
+            self.abort_queue.put(req_id)
+
+    def _process_aborts_queue(self) -> None:
+        if self.abort_queue.empty():
+            return
+
+        request_ids: list[str] = []
+        while not self.abort_queue.empty():
+            ids = self.abort_queue.get_nowait()
+            request_ids.extend((ids,) if isinstance(ids, str) else ids)
+
+        self._abort_requests(request_ids)
+
+    def _abort_requests(self, request_ids: str | Iterable[str]) -> None:
+        request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids)
+
+        sched_req_ids: list[str] = []
+        for request_id in dict.fromkeys(request_ids):
+            sched_req_id = self.scheduler.get_sched_req_id(request_id)
+            if sched_req_id is not None:
+                sched_req_ids.append(sched_req_id)
+
+        for sched_req_id in dict.fromkeys(sched_req_ids):
+            if self.scheduler.get_request_state(sched_req_id) is not None:
+                self.scheduler.finish_requests(sched_req_id, DiffusionRequestStatus.FINISHED_ABORTED)
+
+    def _finalize_finished_request(
+        self,
+        sched_req_id: str,
+        runner_output: RunnerOutput | None = None,
+        missing_result_error: str = "Diffusion scheduler finished target request without execution output.",
+    ) -> DiffusionOutput:
+        state = self.scheduler.get_request_state(sched_req_id)
+        popped_state = self.scheduler.pop_request_state(sched_req_id)
+        state = state or popped_state
+
+        if state is None:
+            raise RuntimeError(f"Diffusion scheduler lost state for request {sched_req_id}.")
+
+        if state.status == DiffusionRequestStatus.FINISHED_ABORTED:
+            request_id = state.req.request_ids[0] if state.req.request_ids else sched_req_id
+            return DiffusionOutput(
+                aborted=True,
+                abort_message=f"Request {request_id} aborted.",
+            )
+
+        if runner_output is not None and runner_output.result is not None:
+            return runner_output.result
+
+        return DiffusionOutput(error=missing_result_error)

vllm_omni/diffusion/distributed/group_coordinator.py

@@ -213,7 +213,7 @@ class GroupCoordinator:
         input_size[0] *= world_size
         output_tensor = torch.empty(input_size, dtype=input_.dtype, device=input_.device)
         # All-gather.
-        torch.distributed.all_gather_into_tensor(output_tensor, input_, group=self.device_group)
+        torch.distributed.all_gather_into_tensor(output_tensor, input_.contiguous(), group=self.device_group)
         if dim != 0:
             input_size[0] //= world_size
             output_tensor = output_tensor.reshape(

vllm_omni/diffusion/executor/abstract.py

@@ -1,11 +1,17 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
+    from vllm_omni.diffusion.worker.utils import RunnerOutput
+
 
 class DiffusionExecutor(ABC):
     """Abstract base class for Diffusion executors."""
@@ -13,7 +19,7 @@ class DiffusionExecutor(ABC):
     uses_multiproc: bool = False
 
     @staticmethod
-    def get_class(od_config: OmniDiffusionConfig) -> type["DiffusionExecutor"]:
+    def get_class(od_config: OmniDiffusionConfig) -> type[DiffusionExecutor]:
         executor_class: type[DiffusionExecutor]
         distributed_executor_backend = od_config.distributed_executor_backend
 
@@ -63,6 +69,16 @@ class DiffusionExecutor(ABC):
         """Add requests to the execution queue."""
         pass
 
+    @abstractmethod
+    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+        """Execute request-mode work from a scheduler output."""
+        pass
+
+    @abstractmethod
+    def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+        """Execute step-mode work from a scheduler output."""
+        pass
+
     @abstractmethod
     def collective_rpc(
         self,
@@ -71,6 +87,7 @@ class DiffusionExecutor(ABC):
         args: tuple = (),
         kwargs: dict | None = None,
         unique_reply_rank: int | None = None,
+        exec_all_ranks: bool = False,
     ) -> Any:
         """Execute a method on workers."""
         pass

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import multiprocessing as mp
 import time
 import weakref
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import zmq
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
@@ -14,6 +16,10 @@ from vllm_omni.diffusion.ipc import unpack_diffusion_output_shm
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker import WorkerProc
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
+    from vllm_omni.diffusion.worker.utils import RunnerOutput
+
 logger = init_logger(__name__)
 
 
@@ -190,6 +196,61 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             logger.error(f"Generate call failed: {e}")
             raise
 
+    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+        """Adapt request-mode scheduler output to worker execute_model RPC."""
+        from vllm_omni.diffusion.worker.utils import RunnerOutput
+
+        self._ensure_open()
+        if scheduler_output.num_scheduled_reqs != 1:
+            raise ValueError(
+                f"Request mode currently supports batch_size=1, "
+                f"but got {scheduler_output.num_scheduled_reqs} scheduled requests."
+            )
+
+        new_req = scheduler_output.scheduled_new_reqs[0]
+        result = self.collective_rpc(
+            "execute_model",
+            args=(new_req.req, self.od_config),
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        if not isinstance(result, DiffusionOutput):
+            raise RuntimeError(f"Unexpected response type for execute_request: {type(result)!r}")
+
+        return RunnerOutput(
+            req_id=new_req.sched_req_id,
+            step_index=None,
+            finished=True,
+            result=result,
+        )
+
+    def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+        """Forward step-mode scheduler output to worker execute_stepwise RPC."""
+        from vllm_omni.diffusion.worker.utils import RunnerOutput
+
+        self._ensure_open()
+        result = self.collective_rpc(
+            "execute_stepwise",
+            args=(scheduler_output,),
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+
+        if isinstance(result, RunnerOutput):
+            return result
+        # TODO: Remove this fallback; DiffusionOutput cannot faithfully represent
+        # failed multi-request step batches.
+        if isinstance(result, DiffusionOutput):
+            req_id = scheduler_output.scheduled_req_ids[0] if scheduler_output.scheduled_req_ids else ""
+            return RunnerOutput(
+                req_id=req_id,
+                step_index=None,
+                finished=True,
+                result=result,
+            )
+        else:
+            raise RuntimeError(f"Unexpected response type for execute_step: {type(result)!r}")
+
     def collective_rpc(
         self,
         method: str,
@@ -197,6 +258,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         args: tuple = (),
         kwargs: dict | None = None,
         unique_reply_rank: int | None = None,
+        exec_all_ranks: bool = False,
     ) -> Any:
         self._ensure_open()
 
@@ -212,7 +274,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             "args": args,
             "kwargs": kwargs,
             "output_rank": unique_reply_rank if unique_reply_rank is not None else 0,
-            "exec_all_ranks": unique_reply_rank is None,
+            "exec_all_ranks": unique_reply_rank is None or exec_all_ranks,
         }
 
         try:
@@ -228,6 +290,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 try:
                     response = self._result_mq.dequeue(timeout=dequeue_timeout)
 
+                    try:
+                        unpack_diffusion_output_shm(response)
+                    except Exception as e:
+                        logger.warning("SHM unpack failed (data may already be inline): %s", e)
+
                     # Check if response indicates an error
                     if isinstance(response, dict) and response.get("status") == "error":
                         raise RuntimeError(

vllm_omni/diffusion/layers/custom_op.py

@@ -25,6 +25,8 @@ class CustomOp(nn.Module):
             return self.forward_npu
         elif current_omni_platform.is_xpu():
             return self.forward_xpu
+        elif current_omni_platform.is_musa():
+            return self.forward_musa
         else:
             return self.forward_native
 
@@ -51,3 +53,7 @@ class CustomOp(nn.Module):
     def forward_hip(self, *args, **kwargs):
         # By default, we assume that HIP ops are compatible with CUDA ops.
         return self.forward_cuda(*args, **kwargs)
+
+    def forward_musa(self, *args, **kwargs):
+        # By default, we assume that MUSA ops are compatible with CUDA ops.
+        return self.forward_cuda(*args, **kwargs)

vllm_omni/diffusion/lora/manager.py

@@ -218,10 +218,16 @@ class DiffusionLoRAManager:
             lora_scale: The external scale for the LoRA adapter.
         """
         if lora_request is None:
+            if self._active_adapter_id is None:
+                logger.debug("No lora_request provided and adapters are already inactive")
+                return
             logger.debug("No lora_request provided, deactivating all LoRA adapters")
             self._deactivate_all_adapters()
             return
         elif math.isclose(0.0, lora_scale):
+            if self._active_adapter_id is None:
+                logger.debug("Received LoRA scale 0 with adapters already inactive")
+                return
             logger.warning("Received a request with LoRA scale 0; deactivating all LoRA adapters")
             self._deactivate_all_adapters()
             return
@@ -605,6 +611,9 @@ class DiffusionLoRAManager:
         self._update_adapter_scale(adapter_id, scale)
 
     def _deactivate_all_adapters(self) -> None:
+        if self._active_adapter_id is None:
+            logger.debug("All adapters already inactive")
+            return
         logger.info("Deactivating all adapters: %d layers", len(self._lora_modules))
         for lora_layer in self._lora_modules.values():
             lora_layer.reset_lora(0)

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -387,7 +387,12 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
 
         else:
             image_input = (
-                None if isinstance(first_prompt, str) else (first_prompt.get("multi_modal_data") or {}).get("image")
+                None
+                if isinstance(first_prompt, str)
+                else (
+                    (first_prompt.get("multi_modal_data") or {}).get("image")
+                    or (first_prompt.get("multi_modal_data") or {}).get("img2img")
+                )
             )
             if image_input and not isinstance(image_input, list):
                 image_input = [image_input]

vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py

@@ -61,6 +61,9 @@ from vllm_omni.diffusion.attention.backends.abstract import (
 )
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
     get_pp_group,
     get_sequence_parallel_rank,
     get_sequence_parallel_world_size,
@@ -2535,6 +2538,61 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
     def set_scheduler(self, new_scheduler):
         self.register_modules(scheduler=new_scheduler)
 
+    @staticmethod
+    def _split_model_kwargs_for_cfg_parallel(model_kwargs: dict[str, Any], batch_size: int, cfg_rank: int) -> None:
+        """Split batch-doubled model_kwargs in-place for CFG parallel.
+
+        The tokenizer produces inputs with cfg_factor=2, so all batch-dim
+        tensors have shape [batch_size*2, ...]. This method slices them
+        so that rank 0 gets the conditioned half and rank 1 gets the
+        unconditioned half.
+        """
+        s = slice(cfg_rank * batch_size, (cfg_rank + 1) * batch_size)
+
+        # Tensor fields with leading batch dimension
+        tensor_keys = [
+            "position_ids",
+            "image_mask",
+            "gen_timestep_scatter_index",
+            "cond_vae_image_mask",
+            "cond_vit_image_mask",
+            "cond_timestep_scatter_index",
+        ]
+        for key in tensor_keys:
+            if key in model_kwargs and model_kwargs[key] is not None:
+                model_kwargs[key] = model_kwargs[key][s]
+
+        # custom_pos_emb: tuple of (cos, sin)
+        if "custom_pos_emb" in model_kwargs and model_kwargs["custom_pos_emb"] is not None:
+            cos, sin = model_kwargs["custom_pos_emb"]
+            model_kwargs["custom_pos_emb"] = (cos[s], sin[s])
+
+        # cond_vae_images: tensor or list
+        if model_kwargs.get("cond_vae_images") is not None:
+            v = model_kwargs["cond_vae_images"]
+            if isinstance(v, torch.Tensor):
+                model_kwargs["cond_vae_images"] = v[s]
+            elif isinstance(v, list):
+                model_kwargs["cond_vae_images"] = v[s.start : s.stop]
+
+        # cond_timestep: tensor or list
+        if model_kwargs.get("cond_timestep") is not None:
+            v = model_kwargs["cond_timestep"]
+            if isinstance(v, torch.Tensor):
+                model_kwargs["cond_timestep"] = v[s]
+            elif isinstance(v, list):
+                model_kwargs["cond_timestep"] = v[s.start : s.stop]
+
+        # cond_vit_images: list of tensors
+        if model_kwargs.get("cond_vit_images") is not None:
+            model_kwargs["cond_vit_images"] = model_kwargs["cond_vit_images"][s.start : s.stop]
+
+        # vit_kwargs: dict of lists
+        if model_kwargs.get("vit_kwargs") is not None:
+            model_kwargs["vit_kwargs"] = {
+                k: v[s.start : s.stop] if isinstance(v, list) else v[s] for k, v in model_kwargs["vit_kwargs"].items()
+            }
+
     @torch.no_grad()
     def __call__(
         self,
@@ -2621,7 +2679,8 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         self._guidance_scale = guidance_scale
         self._guidance_rescale = guidance_rescale
 
-        cfg_factor = 1 + self.do_classifier_free_guidance
+        # Detect CFG parallel configuration (only 2-branch layout is supported)
+        cfg_parallel_ready = self.do_classifier_free_guidance and get_classifier_free_guidance_world_size() == 2
 
         # Define call parameters
         device = self._execution_device
@@ -2649,13 +2708,33 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         # Prepare extra step kwargs.
         _scheduler_step_extra_kwargs = self.prepare_extra_func_kwargs(self.scheduler.step, {"generator": generator})
 
-        # Prepare model kwargs
+        # Prepare model kwargs — attention mask is built from the full
+        # (cfg_factor=2) batch before any splitting so that each rank's
+        # slice is correct.
         input_ids = model_kwargs.pop("input_ids")
         attention_mask = self.model._prepare_attention_mask_for_generation(  # noqa
             input_ids,
             self.model.generation_config,
             model_kwargs=model_kwargs,
         )
+
+        # Split inputs for CFG parallel: each rank processes only its branch.
+        if cfg_parallel_ready:
+            cfg_group = get_cfg_group()
+            cfg_rank = get_classifier_free_guidance_rank()
+
+            # Ensure all ranks start with the same latents
+            latents = latents.contiguous()
+            cfg_group.broadcast(latents, src=0)
+
+            # Split batch-doubled tensors: rank 0 → conditioned, rank 1 → unconditioned
+            s = slice(cfg_rank * batch_size, (cfg_rank + 1) * batch_size)
+            input_ids = input_ids[s]
+            attention_mask = attention_mask[s]
+            self._split_model_kwargs_for_cfg_parallel(model_kwargs, batch_size, cfg_rank)
+        else:
+            cfg_factor = 1 + self.do_classifier_free_guidance
+
         b, _, q_len1, seq_len = attention_mask.shape
         query_lens = [q_len1] * b
         seq_lens = [seq_len] * b
@@ -2678,9 +2757,12 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
 
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
-                # expand the latents if we are doing classifier free guidance
-                latent_model_input = torch.cat([latents] * cfg_factor)
-                # latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                if cfg_parallel_ready:
+                    # CFG parallel: each rank forwards its own branch (no batch doubling)
+                    latent_model_input = latents
+                else:
+                    # Sequential CFG: double the batch
+                    latent_model_input = torch.cat([latents] * cfg_factor)
 
                 t_expand = t.repeat(latent_model_input.shape[0])
 
@@ -2721,14 +2803,17 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                     # TeaCache fast path: reuse previous prediction
                     pred = tc_prev_pred
 
-                # perform guidance
-                if self.do_classifier_free_guidance:
+                # Perform guidance
+                if cfg_parallel_ready:
+                    # CFG parallel: all_gather → all ranks combine locally (no broadcast needed)
+                    gathered = cfg_group.all_gather(pred, separate_tensors=True)
+                    pred = self.cfg_operator(gathered[0], gathered[1], self.guidance_scale, step=i)
+                elif self.do_classifier_free_guidance:
                     pred_cond, pred_uncond = pred.chunk(2)
                     pred = self.cfg_operator(pred_cond, pred_uncond, self.guidance_scale, step=i)
 
-                # compute the previous noisy sample x_t -> x_t-1
+                # Scheduler step (all ranks compute locally in CFG parallel)
                 latents = self.scheduler.step(pred, t, latents, **_scheduler_step_extra_kwargs, return_dict=False)[0]
-
                 if i != len(timesteps) - 1 and should_compute:
                     model_kwargs = self.model._update_model_kwargs_for_generation(  # noqa
                         model_output,

vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py

@@ -149,13 +149,12 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin, Diffu
             "time_embed_2",
             "final_layer.model",
         ]
-
-        device_str = f"{get_local_device()}"
+        device = get_local_device()
         named_modules = dict(self.named_modules())
         for prefix in non_model_layer_prefixes:
             mod = named_modules.get(prefix)
             if mod:
-                mod.to(device_str)
+                mod.to(device)
 
         unexpected_keywords = [
             "guidance_emb",

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -6,6 +6,7 @@ from typing import Any
 
 import torch
 from torch import nn
+from torch.distributed.tensor import DTensor
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
@@ -58,6 +59,31 @@ class LayerwiseOffloadHook(ModelHook):
         self.dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {}
         self.dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {}
 
+    @staticmethod
+    def _is_dtensor(t: torch.Tensor) -> bool:
+        return isinstance(t, DTensor)
+
+    @staticmethod
+    def _set_tensor_storage(target: torch.Tensor, value: torch.Tensor) -> None:
+        if LayerwiseOffloadHook._is_dtensor(target):
+            target._local_tensor = value
+        else:
+            target.data = value
+
+    @staticmethod
+    def _make_offload_placeholder(tensor: torch.Tensor) -> torch.Tensor:
+        if LayerwiseOffloadHook._is_dtensor(tensor):
+            local_shape = tuple(tensor.to_local().shape)
+            return torch.empty(local_shape, device="meta", dtype=tensor.dtype)
+        return torch.empty((0,), device=tensor.device, dtype=tensor.dtype)
+
+    @staticmethod
+    def _is_materialized_tensor(t: torch.Tensor) -> bool:
+        if LayerwiseOffloadHook._is_dtensor(t):
+            local_t = t.to_local()
+            return not local_t.is_meta
+        return not t.is_meta and t.data.numel() > 0
+
     def initialize_hook(self, module: nn.Module) -> nn.Module:
         # This all happen during the hook instance being registered to hook registry;
         # the input module is kept intact
@@ -71,7 +97,10 @@ class LayerwiseOffloadHook(ModelHook):
 
         # Pre-allocate gpu tensors in a flattened way
         self.dtype_cpu_flattened_weights, self.dtype_metadata = LayerwiseOffloadHook._to_cpu(
-            self.next_block_parameters, self.next_block_buffers, self.device, self.pin_memory
+            self.next_block_parameters,
+            self.next_block_buffers,
+            self.device,
+            self.pin_memory,
         )
 
         return module
@@ -106,13 +135,17 @@ class LayerwiseOffloadHook(ModelHook):
 
         for dtype, name2weights in dtype_grouped_weights.items():
             # total # of parameters + buffers
-            total_numel = sum(t.numel() for _, t in name2weights.items())
+            weights_with_local = []
+            for name, t in name2weights.items():
+                local_t = t.to_local() if hasattr(t, "to_local") else t
+                weights_with_local.append((name, t, local_t))
+            total_numel = sum(local.numel() for _, _, local in weights_with_local)
             cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=pin_memory)
 
             current_offset = 0
-            for name, param_or_buf in name2weights.items():
-                numel = param_or_buf.numel()
-                cpu_tensor[current_offset : current_offset + numel].copy_(param_or_buf.flatten())
+            for name, original_tensor, local_tensor in weights_with_local:
+                numel = local_tensor.numel()
+                cpu_tensor[current_offset : current_offset + numel].copy_(local_tensor.flatten())
                 if dtype not in dtype_metadata:
                     dtype_metadata[dtype] = []
                 dtype_metadata[dtype].append(
@@ -120,11 +153,13 @@ class LayerwiseOffloadHook(ModelHook):
                         "name": name,
                         "offset": current_offset,
                         "numel": numel,
-                        "shape": param_or_buf.shape,
+                        "shape": local_tensor.shape,
                     }
                 )
 
-                param_or_buf.data = torch.empty((), device=device, dtype=dtype)
+                LayerwiseOffloadHook._set_tensor_storage(
+                    original_tensor, LayerwiseOffloadHook._make_offload_placeholder(original_tensor)
+                )
                 current_offset += numel
 
             dtype_cpu_flattened_weights[dtype] = cpu_tensor
@@ -135,7 +170,7 @@ class LayerwiseOffloadHook(ModelHook):
     def is_materialized(self) -> bool:
         """Check whether this block's parameters hold real data on device."""
         for param in self.block_parameters.values():
-            return param.data.dim() > 0
+            return LayerwiseOffloadHook._is_materialized_tensor(param)
 
         return True
 
@@ -172,8 +207,9 @@ class LayerwiseOffloadHook(ModelHook):
                     layer_params[target_name] if target_name in layer_params else layer_bufs[target_name]
                 )
 
-                target_param_or_buf.data = gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]].view(
-                    metadata["shape"]
+                LayerwiseOffloadHook._set_tensor_storage(
+                    target_param_or_buf,
+                    gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]].view(metadata["shape"]),
                 )
 
         self._prefetch_done = evt
@@ -191,9 +227,9 @@ class LayerwiseOffloadHook(ModelHook):
 
         # free GPU residency
         for _, param in self.block_parameters.items():
-            param.data = torch.empty((), device=self.device, dtype=param.dtype)
+            LayerwiseOffloadHook._set_tensor_storage(param, LayerwiseOffloadHook._make_offload_placeholder(param))
         for _, buf in self.block_buffers.items():
-            buf.data = torch.empty((), device=self.device, dtype=buf.dtype)
+            LayerwiseOffloadHook._set_tensor_storage(buf, LayerwiseOffloadHook._make_offload_placeholder(buf))
 
     def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
         # if the previous hook was skipped and the weights are not on device,
@@ -311,7 +347,11 @@ class LayerWiseOffloadBackend(OffloadBackend):
             # during the last layer compute of the previous request.
             last_block, first_block = blocks[-1], blocks[0]
             last_hook = apply_block_hook(
-                last_block, first_block, self.device, self.copy_stream, self.config.pin_cpu_memory
+                last_block,
+                first_block,
+                self.device,
+                self.copy_stream,
+                self.config.pin_cpu_memory,
             )
             last_hook.prefetch_layer(non_blocking=False)
 
@@ -319,7 +359,13 @@ class LayerWiseOffloadBackend(OffloadBackend):
             # Register hook for each of blocks
             for i, block in enumerate(blocks[:-1]):
                 next_block = blocks[(i + 1) % num_blocks]
-                hook = apply_block_hook(block, next_block, self.device, self.copy_stream, self.config.pin_cpu_memory)
+                hook = apply_block_hook(
+                    block,
+                    next_block,
+                    self.device,
+                    self.copy_stream,
+                    self.config.pin_cpu_memory,
+                )
                 block_hooks.append(hook)
 
             # NOTE(yuanheng-zhao): We make each hook gets a backward reference to the hook

vllm_omni/diffusion/sched/__init__.py

@@ -10,16 +10,18 @@ from vllm_omni.diffusion.sched.interface import (
     SchedulerInterface,
 )
 from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
+from vllm_omni.diffusion.sched.step_scheduler import StepScheduler
 
 Scheduler = RequestScheduler
 
 __all__ = [
+    "DiffusionRequestStatus",
     "CachedRequestData",
     "DiffusionRequestState",
-    "DiffusionRequestStatus",
     "DiffusionSchedulerOutput",
     "NewRequestData",
-    "RequestScheduler",
-    "Scheduler",
     "SchedulerInterface",
+    "RequestScheduler",
+    "StepScheduler",
+    "Scheduler",
 ]

vllm_omni/diffusion/sched/base_scheduler.py

@@ -5,13 +5,21 @@ from __future__ import annotations
 
 from collections import deque
 
+from vllm.logger import init_logger
+
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
+    CachedRequestData,
     DiffusionRequestState,
     DiffusionRequestStatus,
+    DiffusionSchedulerOutput,
+    NewRequestData,
     SchedulerInterface,
 )
 
+logger = init_logger(__name__)
+
 
 class _BaseScheduler(SchedulerInterface):
     """Shared queue/state bookkeeping for diffusion schedulers."""
@@ -24,8 +32,6 @@ class _BaseScheduler(SchedulerInterface):
         self._waiting: deque[str] = deque()
         self._running: list[str] = []
         self._finished_req_ids: set[str] = set()
-        # The current DiffusionEngine execution mode does not support real
-        # request batching well, so we keep this fixed at 1 for now.
         self._max_batch_size: int = 1
 
     def initialize(self, od_config: OmniDiffusionConfig) -> None:
@@ -36,8 +42,67 @@ class _BaseScheduler(SchedulerInterface):
         self._waiting.clear()
         self._running.clear()
         self._finished_req_ids.clear()
+        # The current DiffusionEngine execution mode does not support real
+        # request batching well, so we keep this fixed at 1 for now.
+        # TODO: Add support for multiple concurrent requests
+        self.max_num_running_reqs = 1
         self._reset_scheduler_state()
 
+    def add_request(self, request: OmniDiffusionRequest) -> str:
+        sched_req_id = self._make_sched_req_id(request)
+        return self._add_request_with_sched_req_id(sched_req_id, request)
+
+    def _add_request_with_sched_req_id(self, sched_req_id: str, request: OmniDiffusionRequest) -> str:
+        state = DiffusionRequestState(sched_req_id=sched_req_id, req=request)
+        self._request_states[sched_req_id] = state
+        self._register_request_ids(request.request_ids, sched_req_id)
+        self._waiting.append(sched_req_id)
+        logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, sched_req_id, len(self._waiting))
+        return sched_req_id
+
+    def schedule(self) -> DiffusionSchedulerOutput:
+        scheduled_new_reqs: list[NewRequestData] = []
+        scheduled_cached_req_ids: list[str] = []
+
+        # First, schedule the RUNNING request(s)
+        for sched_req_id in self._running:
+            state = self._request_states.get(sched_req_id)
+            if state is not None:
+                scheduled_cached_req_ids.append(sched_req_id)
+
+        # Second, schedule WAITING requests while capacity remains.
+        while self._waiting and len(self._running) < self.max_num_running_reqs:
+            sched_req_id = self._waiting[0]
+            state = self._request_states.get(sched_req_id)
+            if state is None:
+                self._waiting.popleft()
+                continue
+            if not self._can_schedule_waiting(state):
+                break
+
+            self._waiting.popleft()
+            was_new_request = state.status == DiffusionRequestStatus.WAITING
+            state.status = DiffusionRequestStatus.RUNNING
+            self._running.append(sched_req_id)
+            if was_new_request:
+                scheduled_new_reqs.append(NewRequestData.from_state(state))
+            else:
+                scheduled_cached_req_ids.append(sched_req_id)
+
+        scheduler_output = DiffusionSchedulerOutput(
+            step_id=self._step_id,
+            scheduled_new_reqs=scheduled_new_reqs,
+            scheduled_cached_reqs=CachedRequestData(sched_req_ids=scheduled_cached_req_ids),
+            finished_req_ids=set(self._finished_req_ids),
+            num_running_reqs=len(self._running),
+            num_waiting_reqs=len(self._waiting),
+        )
+
+        # update after schedule
+        self._step_id += 1
+        self._finished_req_ids.clear()
+        return scheduler_output
+
     def has_requests(self) -> bool:
         return bool(self._waiting or self._running)
 
@@ -121,12 +186,32 @@ class _BaseScheduler(SchedulerInterface):
         self._finished_req_ids |= finished_req_ids
         return finished_req_ids
 
+    def _finalize_update_from_output(
+        self,
+        sched_output: DiffusionSchedulerOutput,
+        statuses: dict[str, DiffusionRequestStatus],
+        errors: dict[str, str | None] | None = None,
+    ) -> set[str]:
+        # A scheduled request may be aborted after schedule() but before
+        # update_from_output() processes the runner output. It is already
+        # marked finished at that point, but we still need to surface its id
+        # in this update so the engine can observe the terminal state.
+        finished_req_ids = {
+            sched_req_id for sched_req_id in sched_output.scheduled_req_ids if sched_req_id in self._finished_req_ids
+        }
+        finished_req_ids |= self._finish_requests(statuses, errors)
+        return finished_req_ids
+
     def _reset_scheduler_state(self) -> None:
         """Reset subclass-owned state during initialize()/close()."""
 
     def _pop_extra_request_state(self, sched_req_id: str) -> None:
         """Remove subclass-owned per-request state before popping request state."""
 
+    def _can_schedule_waiting(self, state: DiffusionRequestState) -> bool:
+        del state
+        return True
+
     def _register_request_ids(self, request_ids: list[str], sched_req_id: str) -> None:
         for request_id in request_ids:
             existing = self._request_id_to_sched_req_id.get(request_id)

vllm_omni/diffusion/sched/interface.py

@@ -8,12 +8,16 @@ import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from functools import cached_property
+from typing import TYPE_CHECKING
 
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.worker.utils import RunnerOutput
+
 logger = init_logger(__name__)
 
 
@@ -141,7 +145,7 @@ class SchedulerInterface(ABC):
         """Run one scheduling cycle."""
 
     @abstractmethod
-    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: DiffusionOutput) -> set[str]:
+    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
         """Update scheduler state from executor output."""
 
     @abstractmethod

vllm_omni/diffusion/sched/request_scheduler.py

@@ -3,103 +3,48 @@
 
 from __future__ import annotations
 
-from vllm.logger import init_logger
+from typing import TYPE_CHECKING
 
-from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.base_scheduler import _BaseScheduler
 from vllm_omni.diffusion.sched.interface import (
-    CachedRequestData,
-    DiffusionRequestState,
     DiffusionRequestStatus,
     DiffusionSchedulerOutput,
-    NewRequestData,
 )
 
-logger = init_logger(__name__)
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.worker.utils import RunnerOutput
 
 
 class RequestScheduler(_BaseScheduler):
     """Diffusion scheduler with vLLM-style waiting/running queues."""
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
-        sched_req_id = self._make_sched_req_id(request)
-        state = DiffusionRequestState(sched_req_id=sched_req_id, req=request)
-        self._request_states[sched_req_id] = state
-        self._register_request_ids(request.request_ids, sched_req_id)
-        self._waiting.append(sched_req_id)
-        logger.debug("Scheduler add_request: %s (waiting=%d)", sched_req_id, len(self._waiting))
-        return sched_req_id
+        return super().add_request(request)
 
     def schedule(self) -> DiffusionSchedulerOutput:
-        scheduled_new_reqs: list[NewRequestData] = []
-        scheduled_cached_req_ids: list[str] = []
+        return super().schedule()
 
-        # First, schedule the RUNNING request(s)
-        for sched_req_id in self._running:
-            state = self._request_states.get(sched_req_id)
-            if state is not None:
-                scheduled_cached_req_ids.append(sched_req_id)
-
-        # Second, schedule WAITING requests while capacity remains.
-        while self._waiting and len(self._running) < self._max_batch_size:
-            sched_req_id = self._waiting.popleft()
-            state = self._request_states.get(sched_req_id)
-            if state is None:
-                continue
-            was_new_request = state.status == DiffusionRequestStatus.WAITING
-            state.status = DiffusionRequestStatus.RUNNING
-            self._running.append(sched_req_id)
-            if was_new_request:
-                scheduled_new_reqs.append(NewRequestData.from_state(state))
-            else:
-                scheduled_cached_req_ids.append(sched_req_id)
-
-        scheduler_output = DiffusionSchedulerOutput(
-            step_id=self._step_id,
-            scheduled_new_reqs=scheduled_new_reqs,
-            scheduled_cached_reqs=CachedRequestData(sched_req_ids=scheduled_cached_req_ids),
-            finished_req_ids=set(self._finished_req_ids),
-            num_running_reqs=len(self._running),
-            num_waiting_reqs=len(self._waiting),
-        )
-
-        self._step_id += 1
-        self._finished_req_ids.clear()
-        return scheduler_output
-
-    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: DiffusionOutput) -> set[str]:
+    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
         scheduled_req_ids = sched_output.scheduled_req_ids
         if not scheduled_req_ids:
             return set()
 
-        # A scheduled request may be aborted after schedule() but before
-        # update_from_output() processes the runner output. It is already
-        # marked finished at that point, but we still need to surface its id
-        # in this update so the engine can observe the terminal state.
-        finished_req_ids = {
-            sched_req_id for sched_req_id in scheduled_req_ids if sched_req_id in self._finished_req_ids
-        }
         terminal_statuses: dict[str, DiffusionRequestStatus] = {}
         terminal_errors: dict[str, str | None] = {}
-        # NOTE: request-mode currently assumes one executor call produces one
-        # DiffusionOutput for the single scheduled request in this cycle.
+        result = output.result
         for sched_req_id in scheduled_req_ids:
             state = self._request_states.get(sched_req_id)
             if state is None or state.is_finished():
                 continue
-            if output.error:
+            if result is None:
                 terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
-                terminal_errors[sched_req_id] = output.error
+                terminal_errors[sched_req_id] = "No output result"
+            elif result.error:
+                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[sched_req_id] = result.error
             else:
                 terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_COMPLETED
                 terminal_errors[sched_req_id] = None
 
-        finished_req_ids |= self._finish_requests(terminal_statuses, terminal_errors)
-        return finished_req_ids
-
-    def abort_request(self, sched_req_id: str) -> bool:
-        if self.get_request_state(sched_req_id) is None:
-            return False
-        self.finish_requests(sched_req_id, DiffusionRequestStatus.FINISHED_ABORTED)
-        return True
+        return self._finalize_update_from_output(sched_output, terminal_statuses, terminal_errors)

vllm_omni/diffusion/sched/step_scheduler.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched.base_scheduler import _BaseScheduler
+from vllm_omni.diffusion.sched.interface import (
+    DiffusionRequestStatus,
+    DiffusionSchedulerOutput,
+)
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.worker.utils import RunnerOutput
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class _StepProgress:
+    current_step: int
+    total_steps: int
+
+
+class StepScheduler(_BaseScheduler):
+    """Placeholder scheduler that advances a request one denoise step per update."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._request_progress: dict[str, _StepProgress] = {}
+
+    def _reset_scheduler_state(self) -> None:
+        self._request_progress.clear()
+
+    def add_request(self, request: OmniDiffusionRequest) -> str:
+        sched_req_id = self._make_sched_req_id(request)
+        total_steps = self._get_total_steps(request)
+        if total_steps <= 0:
+            raise ValueError(f"Diffusion request {sched_req_id} must have positive total_steps, got {total_steps}")
+
+        current_step = request.sampling_params.step_index or 0
+        if current_step < 0 or current_step >= total_steps:
+            raise ValueError(
+                f"Diffusion request {sched_req_id} has invalid initial step_index {current_step} "
+                f"for total_steps={total_steps}"
+            )
+
+        request.sampling_params.step_index = current_step
+        sched_req_id = self._add_request_with_sched_req_id(sched_req_id, request)
+        self._request_progress[sched_req_id] = _StepProgress(current_step=current_step, total_steps=total_steps)
+        logger.debug(
+            "StepScheduler add_request: %s (step=%d/%d, waiting=%d)",
+            sched_req_id,
+            current_step,
+            total_steps,
+            len(self._waiting),
+        )
+        return sched_req_id
+
+    def schedule(self) -> DiffusionSchedulerOutput:
+        return super().schedule()
+
+    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
+        scheduled_req_ids = sched_output.scheduled_req_ids
+        if not scheduled_req_ids:
+            return set()
+
+        terminal_statuses: dict[str, DiffusionRequestStatus] = {}
+        terminal_errors: dict[str, str | None] = {}
+        output_error = output.result.error if output.result is not None else None
+        for sched_req_id in scheduled_req_ids:
+            state = self._request_states.get(sched_req_id)
+            progress = self._request_progress.get(sched_req_id)
+            if state is None or progress is None or state.is_finished():
+                continue
+
+            if output_error is not None:
+                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[sched_req_id] = output_error
+                continue
+
+            if output.step_index is None:
+                logger.warning(
+                    "Received RunnerOutput with no step_index for request %s, treating as error",
+                    sched_req_id,
+                )
+                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
+                terminal_errors[sched_req_id] = "Missing step_index in RunnerOutput"
+                continue
+
+            # We assume that the decoding stage is executed immediately after the denoising stage completes.
+            progress.current_step = output.step_index
+            state.req.sampling_params.step_index = output.step_index
+            if output.finished:
+                terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_COMPLETED
+                terminal_errors[sched_req_id] = None
+            else:
+                state.error = None
+
+        return self._finalize_update_from_output(sched_output, terminal_statuses, terminal_errors)
+
+    def _pop_extra_request_state(self, sched_req_id: str) -> None:
+        self._request_progress.pop(sched_req_id, None)
+
+    def _get_total_steps(self, request: OmniDiffusionRequest) -> int:
+        sampling = request.sampling_params
+
+        if sampling.timesteps is not None:
+            return self._sequence_length(sampling.timesteps)
+        if sampling.sigmas is not None:
+            return len(sampling.sigmas)
+        return int(sampling.num_inference_steps)
+
+    @staticmethod
+    def _sequence_length(values: Any) -> int:
+        ndim = getattr(values, "ndim", None)
+        if ndim == 0:
+            return 1
+
+        shape = getattr(values, "shape", None)
+        if shape is not None:
+            return int(shape[0])
+
+        return len(values)

vllm_omni/diffusion/stage_diffusion_client.py

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any
 
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.data import DiffusionRequestAbortedError
 from vllm_omni.engine.stage_init_utils import StageMetadata
 from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
 from vllm_omni.outputs import OmniRequestOutput
@@ -74,6 +75,20 @@ class StageDiffusionClient:
         try:
             result = await self._engine.generate(prompt, sampling_params, request_id)
             await self._output_queue.put(result)
+        except asyncio.CancelledError:
+            logger.info(
+                "[StageDiffusionClient] Stage-%s req=%s cancelled",
+                self.stage_id,
+                request_id,
+            )
+            raise
+        except DiffusionRequestAbortedError as e:
+            logger.info(
+                "[StageDiffusionClient] Stage-%s req=%s aborted: %s",
+                self.stage_id,
+                request_id,
+                e,
+            )
         except Exception as e:
             logger.exception(
                 "[StageDiffusionClient] Stage-%s req=%s failed: %s",
@@ -138,6 +153,7 @@ class StageDiffusionClient:
             task = self._tasks.pop(rid, None)
             if task:
                 task.cancel()
+        await self._engine.abort(request_ids)
 
     async def collective_rpc_async(
         self,

vllm_omni/diffusion/worker/__init__.py

@@ -2,14 +2,31 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Worker classes for diffusion models."""
 
-from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
-from vllm_omni.diffusion.worker.diffusion_worker import (
-    DiffusionWorker,
-    WorkerProc,
-)
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker, WorkerProc
 
 __all__ = [
     "DiffusionModelRunner",
     "DiffusionWorker",
     "WorkerProc",
 ]
+
+
+def __getattr__(name: str) -> Any:
+    if name == "DiffusionModelRunner":
+        from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+
+        return DiffusionModelRunner
+    if name in {"DiffusionWorker", "WorkerProc"}:
+        from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker, WorkerProc
+
+        return {
+            "DiffusionWorker": DiffusionWorker,
+            "WorkerProc": WorkerProc,
+        }[name]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

vllm_omni/engine/async_omni_engine.py

@@ -860,6 +860,7 @@ class AsyncOmniEngine:
                     "max_num_seqs": 1,
                     "parallel_config": parallel_config,
                     "model_class_name": kwargs.get("model_class_name", None),
+                    "step_execution": kwargs.get("step_execution", False),
                     "vae_use_slicing": kwargs.get("vae_use_slicing", False),
                     "vae_use_tiling": kwargs.get("vae_use_tiling", False),
                     "cache_backend": cache_backend,

vllm_omni/engine/orchestrator.py

@@ -268,6 +268,9 @@ class Orchestrator:
                     continue
                 idle = False
 
+                # Handle prefill-finished KV-ready signals before finished outputs.
+                await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs)
+
                 # 2) Process raw outputs through the output processor
                 request_outputs = await self._process_stage_outputs(stage_id, raw_outputs)
 
@@ -313,25 +316,7 @@ class Orchestrator:
         # CFG companion handling: companions don't produce user-visible output
         # and don't forward to the next stage directly.
         if finished and req_id in self._companion_ids:
-            parent_id = self._companion_to_parent.get(req_id)
-            if parent_id is not None:
-                self._companion_done.setdefault(parent_id, set()).add(req_id)
-                logger.debug(
-                    "[Orchestrator] CFG companion %s done (parent=%s)",
-                    req_id,
-                    parent_id,
-                )
-                # Check if parent is waiting and all companions are done
-                if parent_id in self._deferred_parents and self._all_companions_done(parent_id):
-                    deferred = self._deferred_parents.pop(parent_id)
-                    parent_state = self.request_states.get(parent_id)
-                    if parent_state is not None:
-                        await self._forward_to_next_stage(
-                            parent_id,
-                            deferred["stage_id"],
-                            deferred["output"],
-                            parent_state,
-                        )
+            await self._handle_cfg_companion_ready(req_id)
             self.request_states.pop(req_id, None)
             return
 
@@ -358,17 +343,17 @@ class Orchestrator:
                 }
             )
 
-        if finished and stage_id < req_state.final_stage_id and not self.async_chunk:
-            # If this parent has CFG companions, defer forwarding until all done
+        if (
+            finished
+            and stage_id < req_state.final_stage_id
+            and not self.async_chunk
+            and not self._next_stage_already_submitted(stage_id, req_state)
+        ):
             if req_id in self._companion_map and not self._all_companions_done(req_id):
                 self._deferred_parents[req_id] = {
                     "stage_id": stage_id,
                     "output": output,
                 }
-                logger.debug(
-                    "[Orchestrator] Parent %s deferred, waiting for CFG companions",
-                    req_id,
-                )
             else:
                 await self._forward_to_next_stage(req_id, stage_id, output, req_state)
 
@@ -393,6 +378,56 @@ class Orchestrator:
         done_set = self._companion_done.get(parent_id, set())
         return all(cid in done_set for cid in role_map.values())
 
+    def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool:
+        return (stage_id + 1) in req_state.stage_submit_ts
+
+    async def _handle_cfg_companion_ready(self, req_id: str) -> None:
+        """Mark a CFG companion as done; if all companions are done, flush deferred parent."""
+        parent_id = self._companion_to_parent.get(req_id)
+        if parent_id is None:
+            return
+        done_set = self._companion_done.setdefault(parent_id, set())
+        if req_id in done_set:
+            return
+        done_set.add(req_id)
+        if parent_id in self._deferred_parents and self._all_companions_done(parent_id):
+            deferred = self._deferred_parents.pop(parent_id)
+            parent_state = self.request_states.get(parent_id)
+            if parent_state is not None and not self._next_stage_already_submitted(deferred["stage_id"], parent_state):
+                await self._forward_to_next_stage(
+                    parent_id,
+                    deferred["stage_id"],
+                    deferred["output"],
+                    parent_state,
+                )
+
+    async def _handle_kv_ready_raw_outputs(self, stage_id: int, raw_outputs: EngineCoreOutputs) -> None:
+        """Forward split requests once stage-0 KV is ready, not only when decode fully finishes."""
+        if self.async_chunk:
+            return
+        for raw_output in raw_outputs.outputs:
+            kv_params = getattr(raw_output, "kv_transfer_params", None)
+            if not (isinstance(kv_params, dict) and kv_params.get("kv_ready")):
+                continue
+            req_id = raw_output.request_id
+            req_state = self.request_states.get(req_id)
+            if req_state is None:
+                continue
+            if req_id in self._companion_ids:
+                await self._handle_cfg_companion_ready(req_id)
+                continue
+            if stage_id >= req_state.final_stage_id:
+                continue
+            if self._next_stage_already_submitted(stage_id, req_state):
+                continue
+            if req_id in self._companion_map and not self._all_companions_done(req_id):
+                self._deferred_parents[req_id] = {
+                    "stage_id": stage_id,
+                    "output": raw_output,
+                }
+            else:
+                await self._forward_to_next_stage(req_id, stage_id, raw_output, req_state)
+
     def _build_stage_metrics(
         self,
         stage_id: int,

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -18,7 +18,11 @@ from typing import Any
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_file_to_dict
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
+from vllm_omni.diffusion.data import (
+    DiffusionRequestAbortedError,
+    OmniDiffusionConfig,
+    TransformerConfig,
+)
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
@@ -308,6 +312,11 @@ class AsyncOmniDiffusion:
                 request,
             )
             result = result[0]
+        except asyncio.CancelledError:
+            self.engine.abort(request_id)
+            raise
+        except DiffusionRequestAbortedError:
+            raise
         except Exception as e:
             logger.error("Generation failed for request %s: %s", request_id, e)
             raise RuntimeError(f"Diffusion generation failed: {e}") from e

vllm_omni/entrypoints/cli/serve.py

@@ -267,6 +267,11 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable cache-dit summary logging after diffusion forward passes.",
         )
+        omni_config_group.add_argument(
+            "--step-execution",
+            action="store_true",
+            help="Enable per-step diffusion execution so running requests can be aborted between denoise steps.",
+        )
 
         # VAE memory optimization parameters
         omni_config_group.add_argument(

vllm_omni/entrypoints/openai/api_server.py

@@ -20,7 +20,7 @@ from typing import Annotated, Any, Literal, cast
 import httpx
 import vllm.envs as envs
 from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket
-from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
+from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse
 from PIL import Image
 from pydantic import BaseModel, Field
 from starlette.datastructures import State
@@ -1956,16 +1956,10 @@ async def _run_video_generation_job(
         raise
 
 
-@router.post(
-    "/v1/videos",
-    responses={
-        HTTPStatus.OK.value: {"model": VideoResponse},
-        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
-        HTTPStatus.SERVICE_UNAVAILABLE.value: {"model": ErrorResponse},
-        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
-    },
-)
-async def create_video(
+VIDEO_SYNC_TIMEOUT_S = 600.0
+
+
+async def _parse_video_form(
     raw_request: Request,
     prompt: str = Form(...),
     input_reference: UploadFile | None = File(default=None),
@@ -1988,48 +1982,15 @@ async def create_video(
     negative_prompt: str | None = Form(default=None),
     lora: str | None = Form(default=None),
     extra_params: str | None = Form(default=None),
-) -> VideoResponse:
-    """Create an asynchronous video generation job.
+) -> tuple[VideoGenerationRequest, "OmniOpenAIServingVideo", str, ReferenceImage | None]:
+    """FastAPI dependency that parses video form data, validates inputs,
+    resolves the handler, and decodes any reference image.
 
-    This OpenAI-style endpoint accepts multipart form-data, validates the
-    request payload, persists a queued job record, and starts generation in the
-    background. The response contains metadata for polling job status rather
-    than the generated video bytes.
-
-    Args:
-        raw_request: Raw FastAPI request for accessing app state.
-        prompt: Text prompt describing the requested video.
-        input_reference: Optional uploaded reference image file.
-        image_reference: Optional JSON-encoded reference image descriptor.
-        model: Optional model name supplied by the client.
-        seconds: Optional target duration string accepted by the video API.
-        size: Optional output size string such as ``1280x720``.
-        user: Optional user identifier forwarded in the stored request.
-        width: Optional explicit output width override.
-        height: Optional explicit output height override.
-        num_frames: Optional explicit frame count override.
-        fps: Optional explicit frame rate override.
-        num_inference_steps: Optional inference step override.
-        guidance_scale: Optional primary guidance scale override.
-        guidance_scale_2: Optional secondary guidance scale override.
-        boundary_ratio: Optional boundary ratio override.
-        flow_shift: Optional flow shift override.
-        true_cfg_scale: Optional true CFG scale override.
-        seed: Optional random seed override.
-        negative_prompt: Optional negative prompt.
-        lora: Optional JSON-encoded per-request LoRA configuration.
-        extra_params: Optional model-specific parameters passed directly to the model's extra_args.
-
-    Returns:
-        A queued ``VideoResponse`` that includes the generated job identifier
-        and initial metadata for later retrieval.
-
-    Raises:
-        HTTPException: If the request is invalid, the video handler is
-        unavailable, or job initialization fails.
+    Used by both ``POST /v1/videos`` (async) and ``POST /v1/videos/sync``.
     """
     input_reference_bytes = await input_reference.read() if input_reference is not None else None
     parsed_image_reference = _parse_form_json(image_reference)
+
     if parsed_image_reference is not None and input_reference_bytes is not None:
         raise HTTPException(
             status_code=HTTPStatus.BAD_REQUEST.value,
@@ -2058,7 +2019,6 @@ async def create_video(
         "lora": _parse_form_json(lora, expected_type=dict),
         "extra_params": _parse_form_json(extra_params, expected_type=dict),
     }
-
     request_data = {k: v for k, v in request_data.items() if v is not None}
     request = VideoGenerationRequest(**request_data)
 
@@ -2082,25 +2042,101 @@ async def create_video(
     except HTTPException:
         raise
     except Exception as e:
-        logger.exception("Video generation failed: %s", e)
+        logger.exception("Video generation setup failed: %s", e)
         raise HTTPException(
             status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
-            detail=f"Video generation failed: {str(e)}",
+            detail=f"Video generation setup failed: {str(e)}",
         )
-    ref = video_response_from_request(effective_model_name, request)
 
     try:
         image_data = await decode_input_reference(request.image_reference, input_reference_bytes)
     except InvalidInputReferenceError as exc:
         raise HTTPException(400, detail=str(exc) or "Invalid input reference.") from exc
 
-    reference_image = ReferenceImage(data=image_data) if image_data is not None else image_data
+    reference_image = ReferenceImage(data=image_data) if image_data is not None else None
+    return request, handler, effective_model_name, reference_image
+
+
+@router.post(
+    "/v1/videos",
+    responses={
+        HTTPStatus.OK.value: {"model": VideoResponse},
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.SERVICE_UNAVAILABLE.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+async def create_video(
+    ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form),
+) -> VideoResponse:
+    """Create an asynchronous video generation job.
+
+    Accepts multipart form-data (see ``_parse_video_form`` for parameters),
+    persists a queued job record, and starts generation in the background.
+    """
+    request, handler, effective_model_name, reference_image = ctx
+    ref = video_response_from_request(effective_model_name, request)
     await VIDEO_STORE.upsert(ref.id, ref)
     task = asyncio.create_task(_run_video_generation_job(handler, request, ref.id, reference_image))
     await VIDEO_TASKS.upsert(ref.id, task)
     return ref
 
 
+@router.post(
+    "/v1/videos/sync",
+    responses={
+        HTTPStatus.OK.value: {"content": {"video/mp4": {}}},
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.SERVICE_UNAVAILABLE.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+async def create_video_sync(
+    ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form),
+) -> Response:
+    """Synchronous video generation endpoint.
+
+    Accepts the same form parameters as ``POST /v1/videos`` but blocks until
+    generation completes and returns raw video bytes (``video/mp4``) directly.
+    Designed for benchmark and testing scenarios.
+
+    Metadata is returned via response headers ``X-Request-Id``,
+    ``X-Model``, and ``X-Inference-Time-S``.
+    """
+    request, handler, effective_model_name, reference_image = ctx
+    request_id = f"video_sync-{random_uuid()}"
+    started_at = time.perf_counter()
+    try:
+        video_bytes = await asyncio.wait_for(
+            handler.generate_video_bytes(request, request_id, reference_image=reference_image),
+            timeout=VIDEO_SYNC_TIMEOUT_S,
+        )
+    except asyncio.TimeoutError:
+        raise HTTPException(
+            status_code=HTTPStatus.GATEWAY_TIMEOUT.value,
+            detail=f"Video generation timed out after {VIDEO_SYNC_TIMEOUT_S}s.",
+        )
+    except HTTPException:
+        raise
+    except Exception as exc:
+        logger.exception("Sync video generation failed for request_id=%s", request_id)
+        raise HTTPException(
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+            detail=f"Video generation failed: {str(exc)}",
+        ) from exc
+    inference_time_s = time.perf_counter() - started_at
+
+    return Response(
+        content=video_bytes,
+        media_type="video/mp4",
+        headers={
+            "X-Request-Id": request_id,
+            "X-Model": effective_model_name,
+            "X-Inference-Time-S": f"{inference_time_s:.3f}",
+        },
+    )
+
+
 @router.get("/v1/videos", response_model=VideoListResponse)
 async def list_videos(
     after: str | None = None,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -276,6 +276,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             output_modalities if output_modalities is not None else self.engine_client.output_modalities
         )
 
+        num_inference_steps = None
         # Omni multistage image generation: Stage-0 (AR) should receive a clean
         # text prompt (and optional conditioning image/size) so the model's own
         # processor can construct the correct inputs.
@@ -309,6 +310,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     extra_body = request.model_extra or {}
                 height = extra_body.get("height")
                 width = extra_body.get("width")
+                num_inference_steps = extra_body.get("num_inference_steps")
+                if num_inference_steps is not None:
+                    try:
+                        num_inference_steps = int(num_inference_steps)
+                    except Exception:
+                        num_inference_steps = None
                 if "size" in extra_body:
                     try:
                         size_str = extra_body["size"]
@@ -372,14 +379,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     # Use standard OpenAI API parameters for comprehension stage
                     sampling_params_list = self._build_sampling_params_list_from_request(request)
 
-                # Apply user-specified height/width to diffusion stage(s) for image generation
-                if _image_gen_height is not None or _image_gen_width is not None:
+                # Apply user-specified overrides to diffusion stage(s) for image generation
+                if _image_gen_height is not None or _image_gen_width is not None or num_inference_steps is not None:
                     for idx, sp in enumerate(sampling_params_list):
-                        # Diffusion stages typically have height/width attributes
                         if hasattr(sp, "height") and _image_gen_height is not None:
                             sp.height = _image_gen_height
                         if hasattr(sp, "width") and _image_gen_width is not None:
                             sp.width = _image_gen_width
+                        if hasattr(sp, "num_inference_steps") and num_inference_steps is not None:
+                            sp.num_inference_steps = num_inference_steps
 
                 self._log_inputs(
                     request_id,

vllm_omni/entrypoints/openai/serving_video.py

@@ -20,7 +20,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoGenerationResponse,
 )
 from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
-from vllm_omni.entrypoints.openai.video_api_utils import encode_video_base64
+from vllm_omni.entrypoints.openai.video_api_utils import _encode_video_bytes, encode_video_base64
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams, OmniTextPrompt
 
 logger = init_logger(__name__)
@@ -71,13 +71,18 @@ class OmniOpenAIServingVideo:
             stage_configs=stage_configs,
         )
 
-    async def generate_videos(
+    async def _run_and_extract(
         self,
         request: VideoGenerationRequest,
         reference_id: str,
         *,
         reference_image: ReferenceImage | None = None,
-    ) -> VideoGenerationResponse:
+    ) -> tuple[list[Any], list[Any | None], int, int]:
+        """Run the generation pipeline and extract video/audio outputs.
+
+        Returns:
+            Tuple of (videos, audios, audio_sample_rate, output_fps).
+        """
         prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt)
         if request.negative_prompt is not None:
             prompt["negative_prompt"] = request.negative_prompt
@@ -144,12 +149,23 @@ class OmniOpenAIServingVideo:
         )
 
         result = await self._run_generation(prompt, gen_params, reference_id)
-        _t_encode_start = time.perf_counter()
         videos = self._extract_video_outputs(result)
         audios = self._extract_audio_outputs(result, expected_count=len(videos))
         audio_sample_rate = self._resolve_audio_sample_rate(result)
         output_fps = vp.fps or 24
+        return videos, audios, audio_sample_rate, output_fps
 
+    async def generate_videos(
+        self,
+        request: VideoGenerationRequest,
+        reference_id: str,
+        *,
+        reference_image: ReferenceImage | None = None,
+    ) -> VideoGenerationResponse:
+        videos, audios, audio_sample_rate, output_fps = await self._run_and_extract(
+            request, reference_id, reference_image=reference_image
+        )
+        _t_encode_start = time.perf_counter()
         video_data = [
             VideoData(
                 b64_json=(
@@ -169,6 +185,32 @@ class OmniOpenAIServingVideo:
         logger.info("Video response encoding (MP4+base64): %.2f ms", _t_encode_ms)
         return VideoGenerationResponse(created=int(time.time()), data=video_data)
 
+    async def generate_video_bytes(
+        self,
+        request: VideoGenerationRequest,
+        reference_id: str,
+        *,
+        reference_image: ReferenceImage | None = None,
+    ) -> bytes:
+        """Generate a video and return raw MP4 bytes, bypassing base64 encoding."""
+        videos, audios, audio_sample_rate, output_fps = await self._run_and_extract(
+            request, reference_id, reference_image=reference_image
+        )
+        if len(videos) > 1:
+            logger.warning(
+                "Video request %s generated %d outputs; returning only the first.", reference_id, len(videos)
+            )
+        audio = audios[0]
+        _t_encode_start = time.perf_counter()
+        video_bytes = _encode_video_bytes(
+            videos[0],
+            fps=output_fps,
+            **({"audio": audio, "audio_sample_rate": audio_sample_rate} if audio is not None else {}),
+        )
+        _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000
+        logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms)
+        return video_bytes
+
     @staticmethod
     def _apply_lora(lora_body: Any, gen_params: OmniDiffusionSamplingParams) -> None:
         try:

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -1122,7 +1122,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         )
         # Prefer GPU for encoder if available; otherwise keep CPU.
         dev = next(self.parameters()).device
-        if getattr(dev, "type", None) == "cuda":
+        if dev.type != "cpu":
             try:
                 tok.model.to(dev)
                 tok.device = dev
@@ -1450,10 +1450,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             )
             if not speaker:
                 raise ValueError("CustomVoice requires additional_information.speaker.")
-            spk_id_map = getattr(self.talker_config, "spk_id", None) or {}
+            spk_id_map = {k.lower(): v for k, v in (getattr(self.talker_config, "spk_id", None) or {}).items()}
             if speaker not in spk_id_map:
                 raise ValueError(f"Unsupported speaker: {speaker}")
-            spk_id = spk_id_map[speaker.lower()]
+            spk_id = spk_id_map[speaker]
             # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
             spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long)
             spk_embed = self.embed_input_ids(spk_tensor)

vllm_omni/model_executor/stage_configs/bagel_single_stage.yaml

@@ -0,0 +1,32 @@
+# Stage 0: Thinker (multimodal understanding + text generation)
+
+stage_args:
+
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: dit
+      max_num_seqs: 1
+      gpu_memory_utilization: 0.45
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: image
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 1
+
+    final_output: true
+    final_output_type: image
+    is_comprehension: false
+    default_sampling_params:
+      seed: 52
+
+# Runtime edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -212,7 +212,7 @@ def talker2code2wav_async_chunk(
             }
         return None
 
-    in_initial_phase = initial_chunk_size > 0 and initial_chunk_size < chunk_size and length < chunk_size
+    in_initial_phase = initial_chunk_size > 0 and initial_chunk_size < chunk_size and length <= chunk_size
 
     if in_initial_phase:
         # IC phase: emit every initial_chunk_size frames with growing left context.
@@ -225,7 +225,7 @@ def talker2code2wav_async_chunk(
         # Normal phase: offset so the first normal emit picks up after IC phase.
         # IC is stateless (may change with load); any mismatch is absorbed by left_context.
         initial_coverage = (
-            ((chunk_size - 1) // initial_chunk_size) * initial_chunk_size if 0 < initial_chunk_size < chunk_size else 0
+            (chunk_size // initial_chunk_size) * initial_chunk_size if 0 < initial_chunk_size < chunk_size else 0
         )
         adjusted = length - initial_coverage
         if not finished and adjusted % chunk_size != 0:

vllm_omni/platforms/__init__.py

@@ -105,11 +105,28 @@ def xpu_omni_platform_plugin() -> str | None:
     return "vllm_omni.platforms.xpu.platform.XPUOmniPlatform" if is_xpu else None
 
 
+def musa_omni_platform_plugin() -> str | None:
+    """Check if MUSA OmniPlatform should be activated."""
+    is_musa = False
+    logger.debug("Checking if MUSA OmniPlatform is available.")
+    try:
+        import torchada
+
+        if torchada.is_musa_platform():
+            is_musa = True
+            logger.debug("Confirmed MUSA OmniPlatform is available.")
+    except Exception as e:
+        logger.debug("MUSA OmniPlatform is not available because: %s", str(e))
+
+    return "vllm_omni.platforms.musa.platform.MUSAOmniPlatform" if is_musa else None
+
+
 builtin_omni_platform_plugins = {
     "cuda": cuda_omni_platform_plugin,
     "rocm": rocm_omni_platform_plugin,
     "npu": npu_omni_platform_plugin,
     "xpu": xpu_omni_platform_plugin,
+    "musa": musa_omni_platform_plugin,
 }
 
 

vllm_omni/platforms/interface.py

@@ -15,6 +15,7 @@ class OmniPlatformEnum(Enum):
     ROCM = "rocm"
     NPU = "npu"
     XPU = "xpu"
+    MUSA = "musa"
     UNSPECIFIED = "unspecified"
 
 
@@ -41,6 +42,9 @@ class OmniPlatform(Platform):
     def is_rocm(self) -> bool:
         return self._omni_enum == OmniPlatformEnum.ROCM
 
+    def is_musa(self) -> bool:
+        return self._omni_enum == OmniPlatformEnum.MUSA
+
     @classmethod
     def get_omni_ar_worker_cls(cls) -> str:
         raise NotImplementedError

vllm_omni/platforms/musa/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.platforms.musa.platform import MUSAOmniPlatform
+
+__all__ = ["MUSAOmniPlatform"]

vllm_omni/platforms/musa/platform.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+from vllm_musa.platform import MUSAPlatformBase
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+
+logger = init_logger(__name__)
+
+
+class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
+    """MUSA/Moore Threads GPU implementation of OmniPlatform.
+
+    Inherits all MUSA-specific implementations from vllm-musa's MUSAPlatformBase,
+    and adds Omni-specific interfaces from OmniPlatform.
+    """
+
+    _omni_enum = OmniPlatformEnum.MUSA
+
+    @classmethod
+    def get_omni_ar_worker_cls(cls) -> str:
+        return "vllm_omni.platforms.musa.worker.musa_ar_worker.MUSAARWorker"
+
+    @classmethod
+    def get_omni_generation_worker_cls(cls) -> str:
+        return "vllm_omni.platforms.musa.worker.musa_generation_worker.MUSAGenerationWorker"
+
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/model_executor/stage_configs"
+
+    @classmethod
+    def get_diffusion_model_impl_qualname(cls, op_name: str) -> str:
+        # MUSA uses default implementations for diffusion ops
+        if op_name == "hunyuan_fused_moe":
+            return "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault"
+        return super().get_diffusion_model_impl_qualname(op_name)
+
+    @classmethod
+    def prepare_diffusion_op_runtime(cls, op_name: str, **kwargs: Any) -> None:
+        # MUSA uses default runtime preparation
+        return None
+
+    @classmethod
+    def get_diffusion_attn_backend_cls(
+        cls,
+        selected_backend: str | None,
+        head_size: int,
+    ) -> str:
+        """Get the diffusion attention backend class path for MUSA platform.
+
+        MUSA currently supports SDPA (Scaled Dot Product Attention) as the
+        primary backend. Flash Attention support may be added in future
+        when MUSA-specific implementations are available.
+
+        Args:
+            selected_backend: User-selected backend name (e.g., "FLASH_ATTN",
+                "TORCH_SDPA"). If None, uses platform default.
+            head_size: Attention head size.
+
+        Returns:
+            Fully qualified class path of the selected backend.
+        """
+        if selected_backend is not None:
+            backend_upper = selected_backend.upper()
+            backend = DiffusionAttentionBackendEnum[backend_upper]
+            logger.info("Using diffusion attention backend '%s'", backend_upper)
+            return backend.get_path()
+
+        # Default to SDPA for MUSA as it's the most compatible backend
+        logger.info("Defaulting to diffusion attention backend SDPA")
+        return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+    @classmethod
+    def supports_torch_inductor(cls) -> bool:
+        """MUSA supports torch.compile with inductor backend."""
+        return True
+
+    @classmethod
+    def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
+        """Get the torch device for MUSA platform.
+
+        Args:
+            local_rank: Optional local rank for multi-GPU setups.
+
+        Returns:
+            torch.device for MUSA GPU.
+        """
+        if local_rank is None:
+            return torch.device("musa")
+        return torch.device("musa", local_rank)
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        """Get the number of available MUSA devices."""
+        return torch.musa.device_count()
+
+    @classmethod
+    def synchronize(cls) -> None:
+        """Synchronize all MUSA operations."""
+        torch.musa.synchronize()
+
+    @classmethod
+    def get_free_memory(cls, device: torch.device | None = None) -> int:
+        """Get the free memory on the MUSA device.
+
+        Args:
+            device: Optional device to query. If None, uses current device.
+
+        Returns:
+            Free memory in bytes.
+        """
+        free, _ = torch.musa.mem_get_info(device)
+        return free
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return torch.musa.get_device_name(device_id)

vllm_omni/platforms/musa/worker/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.platforms.musa.worker.musa_ar_worker import MUSAARWorker
+from vllm_omni.platforms.musa.worker.musa_generation_worker import (
+    MUSAGenerationWorker,
+)
+
+__all__ = ["MUSAARWorker", "MUSAGenerationWorker"]

vllm_omni/platforms/musa/worker/musa_ar_worker.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""MUSA AR (Autoregressive) worker for vLLM-Omni.
+
+This worker handles autoregressive model stages (thinker/talker) on MUSA devices.
+"""
+
+import gc
+import os
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.mem_utils import MemorySnapshot, format_gib
+from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.utils import report_usage_stats
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+from vllm.v1.worker.utils import request_memory
+from vllm.v1.worker.workspace import init_workspace_manager
+
+from vllm_omni.worker.base import OmniGPUWorkerBase
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+from vllm_omni.worker.mixins import OmniWorkerMixin
+
+logger = init_logger(__name__)
+
+
+class MUSAARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
+    """MUSA AR worker for thinker/talker stages in Omni model."""
+
+    def init_device(self):
+        """Initialize the MUSA device for this worker."""
+        # This env var set by Ray causes exceptions with graph building.
+        os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+        parallel_config = self.parallel_config
+        if (
+            parallel_config.distributed_executor_backend not in ("ray", "external_launcher")
+            and parallel_config.data_parallel_backend != "ray"
+            and parallel_config.nnodes_within_dp == 1
+        ):
+            # Use local DP rank if available, otherwise use global DP rank.
+            dp_local_rank = self.parallel_config.data_parallel_rank_local
+            if dp_local_rank is None:
+                dp_local_rank = self.parallel_config.data_parallel_index
+
+            tp_pp_world_size = self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
+
+            # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
+            self.local_rank += dp_local_rank * tp_pp_world_size
+            assert self.local_rank < torch.musa.device_count(), (
+                f"DP adjusted local rank {self.local_rank} is out of bounds. "
+            )
+            visible_device_count = torch.musa.device_count()
+            assert self.parallel_config.local_world_size <= visible_device_count, (
+                f"local_world_size ({self.parallel_config.local_world_size}) must "
+                f"be less than or equal to the number of visible devices "
+                f"({visible_device_count})."
+            )
+
+        self.device = torch.device(f"musa:{self.local_rank}")
+        torch.musa.set_device(self.device)
+
+        current_platform.check_if_supports_dtype(self.model_config.dtype)
+
+        # Initialize the distributed environment BEFORE taking memory snapshot
+        # This ensures NCCL buffers are allocated before we measure available memory
+        init_worker_distributed_environment(
+            self.vllm_config,
+            self.rank,
+            self.distributed_init_method,
+            self.local_rank,
+            current_platform.dist_backend,
+        )
+
+        # Set random seed.
+        set_random_seed(self.model_config.seed)
+
+        # Now take memory snapshot after distributed environment is initialized
+        gc.collect()
+        torch.musa.empty_cache()
+
+        # Take current memory snapshot
+        self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
+        self.requested_memory = request_memory(init_snapshot, self.cache_config)
+        logger.debug("worker init memory snapshot: %r", self.init_snapshot)
+        logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
+
+        # Initialize workspace manager
+        num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
+        init_workspace_manager(self.device, num_ubatches)
+
+        if self.use_v2_model_runner:
+            # OMNI: v2 model runner does not yet include omni hooks.
+            logger.warning("OMNI MUSAARWorker forces v1 model runner for omni hooks.")
+            self.use_v2_model_runner = False
+
+        # Construct the model runner
+        self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
+
+        if self.rank == 0:
+            # If usage stat is enabled, collect relevant info.
+            report_usage_stats(self.vllm_config)

vllm_omni/platforms/musa/worker/musa_generation_worker.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""MUSA generation worker for vLLM-Omni.
+
+This worker handles non-autoregressive generation stages (e.g., code2wav waveform
+generation) on MUSA devices.
+"""
+
+import gc
+import os
+
+import torch
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.tracing import instrument
+from vllm.utils.mem_utils import MemorySnapshot, format_gib
+from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.utils import report_usage_stats
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+from vllm.v1.worker.utils import request_memory
+from vllm.v1.worker.workspace import init_workspace_manager
+
+from vllm_omni.worker.base import OmniGPUWorkerBase
+from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
+from vllm_omni.worker.mixins import OmniWorkerMixin
+
+logger = init_logger(__name__)
+
+
+class MUSAGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
+    """MUSA generation worker for non-AR waveform generation stage."""
+
+    @instrument(span_name="Init device")
+    def init_device(self):
+        """Initialize the MUSA device for this worker."""
+        # This env var set by Ray causes exceptions with graph building.
+        os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+        parallel_config = self.parallel_config
+        if (
+            parallel_config.distributed_executor_backend not in ("ray", "external_launcher")
+            and parallel_config.data_parallel_backend != "ray"
+            and parallel_config.nnodes_within_dp == 1
+        ):
+            # Use local DP rank if available, otherwise use global DP rank.
+            dp_local_rank = self.parallel_config.data_parallel_rank_local
+            if dp_local_rank is None:
+                dp_local_rank = self.parallel_config.data_parallel_index
+
+            tp_pp_world_size = self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size
+
+            # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
+            self.local_rank += dp_local_rank * tp_pp_world_size
+            assert self.local_rank < torch.musa.device_count(), (
+                f"DP adjusted local rank {self.local_rank} is out of bounds. "
+            )
+            visible_device_count = torch.musa.device_count()
+            assert self.parallel_config.local_world_size <= visible_device_count, (
+                f"local_world_size ({self.parallel_config.local_world_size}) must "
+                f"be less than or equal to the number of visible devices "
+                f"({visible_device_count})."
+            )
+
+        self.device = torch.device(f"musa:{self.local_rank}")
+        torch.musa.set_device(self.device)
+
+        current_platform.check_if_supports_dtype(self.model_config.dtype)
+
+        # Initialize the distributed environment BEFORE taking memory snapshot
+        # This ensures NCCL buffers are allocated before we measure available memory
+        init_worker_distributed_environment(
+            self.vllm_config,
+            self.rank,
+            self.distributed_init_method,
+            self.local_rank,
+            current_platform.dist_backend,
+        )
+
+        # Set random seed.
+        set_random_seed(self.model_config.seed)
+
+        # Now take memory snapshot after distributed environment is initialized
+        gc.collect()
+        torch.musa.empty_cache()
+
+        # Take current memory snapshot
+        self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
+        self.requested_memory = request_memory(init_snapshot, self.cache_config)
+        logger.debug("worker init memory snapshot: %r", self.init_snapshot)
+        logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
+
+        # Initialize workspace manager
+        num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
+        init_workspace_manager(self.device, num_ubatches)
+
+        if self.use_v2_model_runner:
+            # OMNI: v2 model runner does not yet include omni hooks.
+            logger.warning("OMNI MUSAGenerationWorker forces v1 model runner for omni hooks.")
+            self.use_v2_model_runner = False
+
+        # Construct the model runner
+        self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device)
+
+        if self.rank == 0:
+            # If usage stat is enabled, collect relevant info.
+            report_usage_stats(self.vllm_config)

vllm_omni/profiler/omni_torch_profiler.py

@@ -18,11 +18,12 @@ from vllm.profiler.wrapper import WorkerProfiler
 logger = init_logger(__name__)
 
 # NPU has its custom profiler
-TorchProfilerActivity = Literal["CPU", "CUDA", "XPU", "NPU"]
+TorchProfilerActivity = Literal["CPU", "CUDA", "XPU", "NPU", "MUSA"]
 TorchProfilerActivityMap = {
     "CPU": torch.profiler.ProfilerActivity.CPU,
     "CUDA": torch.profiler.ProfilerActivity.CUDA,
     "XPU": torch.profiler.ProfilerActivity.XPU,
+    "MUSA": torch.profiler.ProfilerActivity.CUDA,
 }