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"

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/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__)
 
@@ -1065,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)
 
@@ -1089,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,
@@ -1101,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 ""
 

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/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_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/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

@@ -97,21 +97,50 @@ def test_flush_on_finish():
 
 
 _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),
 ]

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