revert doc deletion

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>

.github/CODEOWNERS

@@ -0,0 +1,17 @@
+# CODEOWNERS is used to automatically assign PR reviewers
+# Format: path pattern @username or @team
+# More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
+
+# Default owners (fallback reviewers for all files)
+# later, we can add more reviewers here
+* @hsliuustc0106
+
+# Assign reviewers by directory
+/.github/ @owner1
+/tests/ @owner1
+/docs/ @owner1
+
+# Example: reviewers for specific directories
+# /examples/offline_inference/ @offline-inference-team
+# /vllm_omni/core/ @core-team
+# /vllm_omni/model_executor/ @model-team

.github/ISSUE_TEMPLATE/200-installation.yml

@@ -7,14 +7,14 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/hsliuustc0106/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
 - type: textarea
   attributes:
     label: Your current environment
     description: |
       Please run the following and paste the output below.
       ```sh
-      wget https://raw.githubusercontent.com/hsliuustc0106/vllm-omni/main/vllm_omni/collect_env.py
+      wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/vllm_omni/collect_env.py
       # For security purposes, please feel free to check the contents of collect_env.py before running it.
       python collect_env.py
       ```

.github/ISSUE_TEMPLATE/400-bug-report.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/hsliuustc0106/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
 - type: markdown
   attributes:
     value: |
@@ -24,7 +24,7 @@ body:
     description: |
       Please run the following and paste the output below.
       ```sh
-      wget https://raw.githubusercontent.com/hsliuustc0106/vllm-omni/main/vllm_omni/collect_env.py
+      wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/vllm_omni/collect_env.py
       # For security purposes, please feel free to check the contents of collect_env.py before running it.
       python collect_env.py
       ```

.github/ISSUE_TEMPLATE/500-feature-request.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/hsliuustc0106/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
 - type: textarea
   attributes:
     label: 🚀 The feature, motivation and pitch

.github/ISSUE_TEMPLATE/600-new-model.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/hsliuustc0106/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
 
       #### We also highly recommend you read the contributing guidelines first to understand how to add a new model.
 - type: textarea

.github/ISSUE_TEMPLATE/700-performance-discussion.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/hsliuustc0106/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+      #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
 - type: textarea
   attributes:
     label: Proposal to improve performance
@@ -19,7 +19,7 @@ body:
   attributes:
     label: Report of performance regression
     description: >
-      Please provide detailed description of performance comparison to confirm the regression. You may want to run the benchmark script at https://github.com/hsliuustc0106/vllm-omni/tree/main/tests/benchmarks .
+      Please provide detailed description of performance comparison to confirm the regression. You may want to run the benchmark script at https://github.com/vllm-project/vllm-omni/tree/main/tests/benchmarks .
   validations:
     required: false
 - type: textarea
@@ -35,7 +35,7 @@ body:
     description: |
       Please run the following and paste the output below.
       ```sh
-      wget https://raw.githubusercontent.com/hsliuustc0106/vllm-omni/main/vllm_omni/collect_env.py
+      wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/vllm_omni/collect_env.py
       # For security purposes, please feel free to check the contents of collect_env.py before running it.
       python collect_env.py
       ```

.github/ISSUE_TEMPLATE/750-RFC.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Please take a look at previous [RFCs](https://github.com/hsliuustc0106/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
+      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
 - type: textarea
   attributes:
     label: Motivation.

.github/ISSUE_TEMPLATE/config.yml

@@ -1,5 +1,5 @@
 blank_issues_enabled: false
 contact_links:
   - name: Questions
-    url: https://github.com/hsliuustc0106/vllm-omni/discussions
+    url: https://github.com/vllm-project/vllm-omni/discussions
     about: Ask questions and discuss with other vLLM-omni community members

.github/PULL_REQUEST_TEMPLATE.md

@@ -18,4 +18,4 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 - [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
 </details>
 
-**BEFORE SUBMITTING, PLEASE READ <https://github.com/hsliuustc0106/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
+**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)

.github/scripts/detect_changed_tests.sh

@@ -0,0 +1,89 @@
+#!/bin/bash
+# Script to detect changed files and find corresponding test files
+# Usage: ./detect_changed_tests.sh [base_sha] [head_sha] [event_type]
+#
+# Outputs to GITHUB_OUTPUT:
+#   test_files: Space-separated list of test files to run
+#   has_tests: 'true' if test files found, 'false' otherwise
+
+set -e
+
+BASE_SHA="${1:-}"
+HEAD_SHA="${2:-}"
+EVENT_TYPE="${3:-push}"
+
+# Determine base and head SHA based on event type
+if [ "$EVENT_TYPE" = "pull_request" ]; then
+    BASE_SHA="${BASE_SHA:-$GITHUB_BASE_REF}"
+    HEAD_SHA="${HEAD_SHA:-$GITHUB_SHA}"
+elif [ -z "$BASE_SHA" ] || [ -z "$HEAD_SHA" ]; then
+    # For push events, use git to find the previous commit
+    BASE_SHA="${BASE_SHA:-$(git rev-parse HEAD~1 2>/dev/null || echo HEAD)}"
+    HEAD_SHA="${HEAD_SHA:-HEAD}"
+fi
+
+echo "Event type: $EVENT_TYPE"
+echo "Base SHA: $BASE_SHA"
+echo "Head SHA: $HEAD_SHA"
+
+# Get changed Python files
+CHANGED_FILES=$(git diff --name-only --diff-filter=ACMRT "$BASE_SHA" "$HEAD_SHA" | grep -E '\.(py)$' || true)
+
+if [ -z "$CHANGED_FILES" ]; then
+    echo "No Python files changed"
+    echo "test_files=" >> "$GITHUB_OUTPUT"
+    echo "has_tests=false" >> "$GITHUB_OUTPUT"
+    exit 0
+fi
+
+echo "Changed Python files:"
+echo "$CHANGED_FILES"
+
+# Extract related test files
+TEST_FILES=""
+for file in $CHANGED_FILES; do
+    # If changed file is a test file, add it directly
+    if [[ "$file" == tests/* ]]; then
+        TEST_FILES="$TEST_FILES $file"
+    # If changed file is source code, find corresponding test file
+    elif [[ "$file" == vllm_omni/* ]]; then
+        # Convert vllm_omni/path/to/module.py to possible test file paths
+        REL_PATH=$(echo "$file" | sed 's|^vllm_omni/||' | sed 's|\.py$||')
+        MODULE_NAME=$(basename "$REL_PATH")
+        DIR_PATH=$(dirname "$REL_PATH")
+
+        # Try multiple possible test file naming patterns
+        POSSIBLE_TESTS=(
+            "tests/test_${REL_PATH//\//_}.py"
+            "tests/${REL_PATH//\//_}_test.py"
+        )
+
+        # If multi-level directory, also try directory structure
+        if [ "$DIR_PATH" != "." ]; then
+            POSSIBLE_TESTS+=(
+                "tests/${DIR_PATH}/test_${MODULE_NAME}.py"
+                "tests/${DIR_PATH}/${MODULE_NAME}_test.py"
+            )
+        fi
+
+        for test_file in "${POSSIBLE_TESTS[@]}"; do
+            if [ -f "$test_file" ]; then
+                TEST_FILES="$TEST_FILES $test_file"
+                break
+            fi
+        done
+    fi
+done
+
+# Deduplicate and format
+TEST_FILES=$(echo "$TEST_FILES" | tr ' ' '\n' | grep -v '^$' | sort -u | tr '\n' ' ' | sed 's/^ *//;s/ *$//')
+
+if [ -z "$TEST_FILES" ]; then
+    echo "No related test files found, will run all tests"
+    echo "test_files=" >> "$GITHUB_OUTPUT"
+    echo "has_tests=false" >> "$GITHUB_OUTPUT"
+else
+    echo "Found test files: $TEST_FILES"
+    echo "test_files=$TEST_FILES" >> "$GITHUB_OUTPUT"
+    echo "has_tests=true" >> "$GITHUB_OUTPUT"
+fi

.github/workflows/ci-all-files.yml

@@ -0,0 +1,60 @@
+name: CI All Files
+
+on:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  test-changed:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0 # Need full history to detect changes
+
+      - name: Free disk space
+        run: |
+          # Remove large, preinstalled toolchains we don't need on the runner
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+          # Prune docker to reclaim space (safe on GH runners)
+          docker system prune -af || true
+          df -h
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Prepare environment with uv
+        run: |
+          python -m pip install --upgrade pip uv
+          uv venv --python "${{ matrix.python-version }}" --seed
+          echo "VIRTUAL_ENV=${PWD}/.venv" >> "$GITHUB_ENV"
+          echo "PATH=${PWD}/.venv/bin:$PATH" >> "$GITHUB_ENV"
+          source .venv/bin/activate
+          uv pip install -e ".[dev]"
+
+      - name: Run tests
+        run: |
+          if [ ! -d tests ]; then
+            echo "tests/ directory not found. Skipping pytest run."
+
+
+            exit 0
+          fi
+          TEST_FILES=$(find tests -type f \( -name "test_*.py" -o -name "*_test.py" \))
+          if [ -z "$TEST_FILES" ]; then
+            echo "No pytest-compatible test files detected under tests/. Skipping pytest run."
+            exit 0
+          fi
+          pytest tests/ -v -m "not slow"

.github/workflows/ci.yml

@@ -0,0 +1,124 @@
+name: CI
+
+on:
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+    inputs:
+      test_changed_only:
+        description: 'Only test changed files (true) or all files (false)'
+        required: false
+        default: 'true'
+        type: choice
+        options:
+          - 'true'
+          - 'false'
+      enable_changed_files_detection:
+        description: 'Enable changed files detection feature'
+        required: false
+        default: 'true'
+        type: choice
+        options:
+          - 'true'
+          - 'false'
+  workflow_call:
+    inputs:
+      test_changed_only:
+        description: 'Only test changed files (true) or all files (false)'
+        required: false
+        type: string
+        default: 'true'
+      enable_changed_files_detection:
+        description: 'Enable changed files detection feature'
+        required: false
+        type: string
+        default: 'true'
+
+permissions:
+  contents: read
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Free disk space
+        run: |
+          # Remove large, preinstalled toolchains we don't need on the runner
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/hostedtoolcache/CodeQL
+          # Prune docker to reclaim space (safe on GH runners)
+          docker system prune -af || true
+          df -h
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Prepare environment with uv
+        run: |
+          python -m pip install --upgrade pip uv
+          uv venv --python "${{ matrix.python-version }}" --seed
+          echo "VIRTUAL_ENV=${PWD}/.venv" >> "$GITHUB_ENV"
+          echo "PATH=${PWD}/.venv/bin:${PATH}" >> "$GITHUB_ENV"
+          # Activate env to ensure subsequent commands succeed in this step
+          source .venv/bin/activate
+          uv pip install -e ".[dev]"
+
+      - name: Determine test scope
+        id: test-scope
+        env:
+          ENABLE_CHANGED_FILES_DETECTION: ${{ inputs.enable_changed_files_detection || 'true' }}
+          TEST_CHANGED_ONLY: ${{ inputs.test_changed_only || 'true' }}
+        run: |
+          # Determine if we should test only changed files
+          # If detection is disabled, always test all files
+          if [ "$ENABLE_CHANGED_FILES_DETECTION" != "true" ]; then
+            TEST_CHANGED_ONLY="false"
+          fi
+
+          echo "test_changed_only=$TEST_CHANGED_ONLY" >> "$GITHUB_OUTPUT"
+          echo "enable_detection=$ENABLE_CHANGED_FILES_DETECTION" >> "$GITHUB_OUTPUT"
+          echo "Test scope: $TEST_CHANGED_ONLY"
+          echo "Changed files detection: $ENABLE_CHANGED_FILES_DETECTION"
+
+      - name: Get changed test files
+        id: changed-tests
+        if: steps.test-scope.outputs.test_changed_only == 'true' && steps.test-scope.outputs.enable_detection == 'true'
+        env:
+          GITHUB_BASE_REF: ${{ github.event.pull_request.base.sha || '' }}
+          GITHUB_SHA: ${{ github.sha }}
+        run: |
+          chmod +x .github/scripts/detect_changed_tests.sh
+          .github/scripts/detect_changed_tests.sh \
+            "${{ github.event.pull_request.base.sha || github.event.before || '' }}" \
+            "${{ github.sha }}" \
+            "${{ github.event_name }}"
+
+      - name: Run tests (changed files only)
+        if: steps.test-scope.outputs.test_changed_only == 'true' && steps.changed-tests.outputs.has_tests == 'true'
+        run: |
+          pytest "${{ steps.changed-tests.outputs.test_files }}" -v -m "not slow"
+
+      - name: Run tests (all files)
+        if: steps.test-scope.outputs.test_changed_only == 'false' || steps.changed-tests.outputs.has_tests == 'false'
+        run: |
+          if [ ! -d tests ]; then
+            echo "tests/ directory not found. Skipping pytest run."
+            exit 0
+          fi
+          TEST_FILES=$(find tests -type f \( -name "test_*.py" -o -name "*_test.py" \))
+          if [ -z "$TEST_FILES" ]; then
+            echo "No pytest-compatible test files detected under tests/. Skipping pytest run."
+            exit 0
+          fi
+          pytest tests/ -v -m "not slow"

.github/workflows/pre-commit.yml

@@ -0,0 +1,24 @@
+name: pre-commit
+
+on:
+  pull_request:
+    branches: [main]
+  workflow_dispatch: # for manual trigger
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+permissions:
+  contents: read
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: 'pip'
+      - uses: pre-commit/action@v3.0.1

.pre-commit-config.yaml

@@ -12,18 +12,21 @@ repos:
       - id: trailing-whitespace
         args: ["--markdown-linebreak-ext=md"]
 
-  - repo: https://github.com/PyCQA/isort
-    rev: 5.13.2
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.14.0
     hooks:
-      - id: isort
+    - id: ruff-check
+      args: [--output-format, github, --fix]
+    - id: ruff-format
 
-  - repo: https://github.com/psf/black
-    rev: 24.10.0
+  - repo: https://github.com/crate-ci/typos
+    rev: v1.35.5
     hooks:
-      - id: black
+      - id: typos
+        # only for staged files
 
-  - repo: https://github.com/PyCQA/flake8
-    rev: 7.1.1
+  - repo: https://github.com/rhysd/actionlint
+    rev: v1.7.7
     hooks:
-      - id: flake8
-        args: ["--max-line-length=88", "--extend-ignore=E203,W503"]
+      - id: actionlint
+        files: ^\.github/workflows/.*\.ya?ml$

.readthedocs.yml

@@ -0,0 +1,24 @@
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+version: 2
+
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.12"
+  jobs:
+    post_checkout:
+      - git fetch --unshallow || true
+
+mkdocs:
+  configuration: mkdocs.yml
+  fail_on_warning: true
+
+# Optionally declare the Python requirements required to build your docs
+python:
+  install:
+    - method: pip
+      path: .
+      extra_requirements:
+        - docs

README.md

@@ -85,7 +85,7 @@ uv pip install --no-build-isolation --editable .
 ```
 
 ### Verification for successful installation of vLLM
-Just run the command below. If no error, it demonstrates that the installation is successfull.
+Just run the command below. If no error, it demonstrates that the installation is successful.
 ```bash
 python -c "import vllm._C"
 ```
@@ -95,20 +95,12 @@ Install additional requirements for vllm-omni
 ```bash
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm_omni
-uv pip install -r requirements.txt
+uv pip install -e .
 ```
 
 ## Run examples (Qwen2.5-omni)
 
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen_2_5_omni
-```
-Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run.
-```bash
-bash run.sh
-```
-The output audio is saved in ./output_audio
+Please check the folder of [examples](examples)
 
 ## Further details
 

docs/DOCS_GUIDE.md

@@ -0,0 +1,139 @@
+# Documentation Build Guide
+
+This directory contains the source files for the vLLM-omni documentation.
+
+## Building Documentation Locally
+
+### Prerequisites
+
+Install documentation dependencies:
+
+```bash
+pip install -e ".[docs]"
+```
+
+### Build and Serve Documentation
+
+From the project root:
+
+```bash
+# Serve documentation locally (auto-reload on changes)
+# This starts a local web server at http://127.0.0.1:8000
+mkdocs serve
+
+# Build static site (generates HTML files in site/ directory)
+mkdocs build
+```
+
+When using `mkdocs serve`, the documentation will be automatically available at `http://127.0.0.1:8000`. The server will automatically reload when you make changes to the documentation files.
+
+## Auto-generating API Documentation
+
+The documentation automatically extracts docstrings from the code using mkdocstrings. To ensure your code is documented:
+
+1. Add docstrings to all public classes, functions, and methods
+2. Use Google or NumPy style docstrings (both are supported)
+3. Rebuild the documentation to see changes
+
+Example docstring:
+
+```python
+class OmniLLM:
+    """Main entry point for vLLM-omni inference.
+
+    This class provides a high-level interface for running multi-modal
+    inference with non-autoregressive models.
+
+    Args:
+        model: Model name or path
+        stage_configs: Optional stage configurations
+        **kwargs: Additional arguments passed to the engine
+
+    Example:
+        >>> llm = OmniLLM(model="Qwen/Qwen2.5-Omni")
+        >>> outputs = llm.generate(prompts="Hello")
+    """
+```
+
+## Documentation Structure
+
+```
+docs/
+├── index.md              # Main documentation page
+├── getting_started/      # Getting started guides
+├── architecture/        # Architecture documentation
+├── api/                 # API reference (auto-generated from code)
+├── examples/            # Code examples
+└── stylesheets/         # Custom CSS
+```
+
+## Publishing Documentation
+
+### GitHub Pages (Recommended)
+
+The documentation is automatically deployed to GitHub Pages using GitHub Actions.
+
+1. **Enable GitHub Pages**:
+   - Go to repository `Settings` → `Pages`
+   - Set `Source` to `GitHub Actions`
+   - Save settings
+
+2. **Push changes**:
+   ```bash
+   git push origin main
+   ```
+
+3. **Documentation will be available at**:
+   - `https://vllm-project.github.io/vllm-omni/`
+
+The GitHub Actions workflow (`.github/workflows/docs.yml`) will automatically:
+- Build the documentation when you push to `main` branch
+- Deploy it to GitHub Pages
+- Update the documentation whenever you make changes
+
+
+### Read the Docs (Alternative)
+
+You can also use Read the Docs for hosting:
+
+1. Sign up at https://readthedocs.org/
+2. Import the `vllm-project/vllm-omni` repository
+3. Read the Docs will automatically build using `.readthedocs.yml`
+4. Documentation will be available at: `https://vllm-omni.readthedocs.io/`
+
+## Configuration
+
+The documentation configuration is in `mkdocs.yml` at the project root.
+
+## Tips
+
+- **API Documentation**: API docs are automatically generated using `mkdocs-api-autonav` and `mkdocstrings`
+  - No need to manually create API pages - they're generated automatically
+  - Use `[module.name.ClassName][]` syntax for cross-references in Summary pages
+- **Code Snippets**: Use `--8<-- "path/to/file.py"` for including code snippets
+- **Markdown**: Use Markdown for all documentation (no need for RST)
+- **Material Theme**: Use Material theme features like:
+  - Admonitions: `!!! note`, `!!! warning`, etc.
+  - Code blocks with syntax highlighting
+  - Tabs for organizing content
+  - Math formulas using `pymdownx.arithmatex`
+
+## Troubleshooting
+
+### Documentation not updating
+
+- Make sure you've saved all files
+- If using `mkdocs serve`, it should auto-reload
+- Check for syntax errors in `mkdocs.yml`
+
+### API links not working
+
+- Ensure class names match exactly (case-sensitive)
+- Check that the module is imported correctly
+- Run `mkdocs build --strict` to check for errors
+
+### Build errors
+
+- Check Python version (requires 3.9+)
+- Ensure all dependencies are installed: `pip install -e ".[docs]"`
+- Check `mkdocs.yml` syntax with `mkdocs build --strict`

docs/api/overview.md

@@ -0,0 +1,68 @@
+# Summary
+
+## Entry Points
+
+Main entry points for vLLM-omni inference and serving.
+
+- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
+- [vllm_omni.entrypoints.omni_llm.OmniStageLLM][]
+- [vllm_omni.entrypoints.omni_stage.OmniStage][]
+- [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
+
+## Inputs
+
+Input data structures for multi-modal inputs.
+
+- [vllm_omni.inputs.data.OmniTokenInputs][]
+- [vllm_omni.inputs.data.OmniTokensPrompt][]
+- [vllm_omni.inputs.data.OmniEmbedsPrompt][]
+- [vllm_omni.inputs.parse.parse_singleton_prompt_omni][]
+- [vllm_omni.inputs.preprocess.OmniInputPreprocessor][]
+
+## Engine
+
+Engine classes for offline and online inference.
+
+- [vllm_omni.engine.arg_utils.OmniEngineArgs][]
+- [vllm_omni.engine.processor.OmniProcessor][]
+- [vllm_omni.engine.output_processor.OmniRequestState][]
+- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
+- [vllm_omni.engine.PromptEmbedsPayload][]
+- [vllm_omni.engine.AdditionalInformationPayload][]
+- [vllm_omni.engine.OmniEngineCoreRequest][]
+
+## Core
+
+Core scheduling and caching components.
+
+- [vllm_omni.core.sched.scheduler.OmniScheduler][]
+- [vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler][]
+- [vllm_omni.core.sched.output.OmniNewRequestData][]
+- [vllm_omni.core.dit_cache_manager.DiTCacheManager][]
+
+## Model Executor
+
+Model execution components.
+
+- [vllm_omni.model_executor.models.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen2_5_omni_thinker.Qwen2_5OmniThinkerForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavForConditionalGenerationVLLM][]
+- [vllm_omni.model_executor.layers.mrope.MRotaryEmbedding][]
+- [vllm_omni.model_executor.models.registry.OmniModelRegistry][]
+
+## Configuration
+
+Configuration classes.
+
+- [vllm_omni.config.OmniModelConfig][]
+
+## Workers
+
+Worker classes and model runners for distributed inference.
+
+- [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
+- [vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
+- [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
+- [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
+- [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]

docs/contributing/README.md

@@ -0,0 +1,170 @@
+# Contributing to vLLM-omni
+
+Thank you for your interest in contributing to vLLM-omni! This document provides guidelines and instructions for contributing.
+
+## Getting Started
+
+### Development Environment
+
+vLLM-omni uses `uv` as the environment manager, to create and manage Python environments. Please follow the documentation to install `uv`. After installing `uv`, you can create a new Python environment using the following commands:
+
+```bash
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+```
+
+Install vLLM from source:
+
+```bash
+git clone https://github.com/vllm-project/vllm.git
+cd vllm
+git checkout v0.11.0
+VLLM_USE_PRECOMPILED=1 uv pip install -e .
+```
+
+Install vLLM-omni from source:
+
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm_omni
+uv pip install -e .
+```
+
+!!! tip
+    vLLM-omni is compatible with Python versions 3.10 to 3.12. However, we recommend developing with Python 3.12 to minimize the chance of your local environment clashing with our CI environment.
+
+### Linting
+
+vLLM-omni uses `pre-commit` to lint and format the codebase. See [pre-commit documentation](https://pre-commit.com/#usage) if `pre-commit` is new to you. Setting up `pre-commit` is as easy as:
+
+```bash
+uv pip install pre-commit
+pre-commit install
+```
+
+vLLM-omni's `pre-commit` hooks will now run automatically every time you commit.
+
+!!! tip
+    You can manually run the `pre-commit` hooks using:
+
+    ```bash
+    pre-commit run     # runs on staged files
+    pre-commit run -a  # runs on all files (short for --all-files)
+    ```
+
+### Documentation
+
+MkDocs is a fast, simple and downright gorgeous static site generator that's geared towards building project documentation. Documentation source files are written in Markdown, and configured with a single YAML configuration file, `mkdocs.yml`.
+
+Get started with:
+
+```bash
+uv pip install -e ".[docs]"
+```
+
+MkDocs comes with a built-in dev-server that lets you preview your documentation as you work on it. From the root of the repository, run:
+
+```bash
+mkdocs serve                           # with API ref (~10 minutes)
+API_AUTONAV_EXCLUDE=vllm_omni mkdocs serve  # API ref off (~15 seconds)
+```
+
+Once you see `Serving on http://127.0.0.1:8000/` in the logs, the live preview is ready! Open <http://127.0.0.1:8000/> in your browser to see it.
+
+For additional features and advanced configurations, refer to the:
+
+- [MkDocs documentation](https://www.mkdocs.org/)
+- [Material for MkDocs documentation](https://squidfunk.github.io/mkdocs-material/) (the MkDocs theme we use)
+
+### Testing
+
+vLLM-omni uses `pytest` to test the codebase.
+
+```bash
+# Install the test dependencies used in CI (CUDA only)
+uv pip install -e ".[dev]" --torch-backend=auto
+
+# Install some common test dependencies (hardware agnostic)
+uv pip install pytest pytest-asyncio
+
+# Run all tests
+pytest tests/
+
+# Run tests for a single test file with detailed output
+pytest -s -v tests/test_logger.py
+```
+
+!!! warning
+    Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
+
+## Issues
+
+If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
+
+!!! important
+    If you discover a security vulnerability, please report it by creating a GitHub issue with the `security` label.
+
+## Pull Requests & Code Reviews
+
+Thank you for your contribution to vLLM-omni! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM-omni maintain the code quality and improve the efficiency of the review process.
+
+### DCO and Signed-off-by
+
+When contributing changes to this project, you must agree to the [DCO](https://developercertificate.org/). Commits must include a `Signed-off-by:` header which certifies agreement with the terms of the DCO.
+
+Using `-s` with `git commit` will automatically add this header.
+
+!!! tip
+    You can enable automatic sign-off via your IDE:
+
+    - **PyCharm**: Click on the `Show Commit Options` icon to the right of the `Commit and Push...` button in the `Commit` window. It will bring up a `git` window where you can modify the `Author` and enable `Sign-off commit`.
+    - **VSCode**: Open the Settings editor and enable the `Git: Always Sign Off` (`git.alwaysSignOff`) field.
+
+### PR Title and Classification
+
+Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:
+
+- `[Bugfix]` for bug fixes.
+- `[CI/Build]` for build or continuous integration improvements.
+- `[Doc]` for documentation fixes and improvements.
+- `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
+- `[Frontend]` For changes on the vLLM-omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
+- `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
+- `[Core]` for changes in the core vLLM-omni logic (e.g., `OmniProcessor`, `OmniScheduler`, etc.)
+- `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.
+- `[Misc]` for PRs that do not fit the above categories. Please use this sparingly.
+
+!!! note
+    If the PR spans more than one category, please include all relevant prefixes.
+
+### Code Quality
+
+The PR needs to meet the following code quality standards:
+
+- We adhere to Google Python style guide and Google C++ style guide.
+- Pass all linter checks.
+- The code needs to be well-documented to ensure future contributors can easily understand the code.
+- Include sufficient tests to ensure the project stays correct and robust. This includes both unit tests and integration tests.
+- Please add documentation to `docs/` if the PR modifies the user-facing behaviors of vLLM-omni. It helps vLLM-omni users understand and utilize the new features or changes.
+
+### Notes for Large Changes
+
+Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with `rfc-required` and might not go through the PR.
+
+### What to Expect for the Reviews
+
+The goal of the vLLM-omni team is to be a _transparent reviewing machine_. We would like to make the review process transparent and efficient and make sure no contributor feels confused or frustrated. However, the vLLM-omni team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:
+
+- After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
+- After the PR is assigned, the reviewer will provide status updates every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM-omni team.
+- After the review, the reviewer will put an `action-required` label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
+- Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.
+- Note that not all CI checks will be executed due to limited computational resources. The reviewer will add `ready` label to the PR when the PR is ready to merge or a full CI run is needed.
+
+## Additional Resources
+
+- [Design Documents](design_documents/index.md) - Architecture and design documentation
+
+## Thank You
+
+Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM-omni. All of your contributions help make vLLM-omni a great tool and community for everyone!

docs/contributing/design_documents/api_design_doc.md

@@ -1,13 +1,10 @@
-# vLLM-omni API Documentation
+# API Design Documentation
 
-This directory contains comprehensive API design documentation for all core modules in vLLM-omni. These templates provide a standardized structure for designing and implementing the core, engine, executor, and worker modules.
+This document contains comprehensive API design documentation for all core modules in vLLM-omni. These templates provide a standardized structure for designing and implementing the core, engine, executor, and worker modules.
 
-## 📋 API Design Templates
+## 📋 Module API
 
-### 1. [API Design Template](API_DESIGN_TEMPLATE.md)
-**Master template** for creating API documentation for any module in vLLM-omni. Use this as a starting point for new modules.
-
-### 2. [Core Module API](core_module_api.md)
+### Core Module API
 **Core module** provides fundamental scheduling, caching, and resource management functionality.
 
 **Key Components:**
@@ -16,7 +13,7 @@ This directory contains comprehensive API design documentation for all core modu
 - Resource allocation and coordination
 - Inter-module communication
 
-### 3. [Engine Module API](engine_module_api.md)
+### Engine Module API
 **Engine module** handles model loading, inference execution, and output processing.
 
 **Key Components:**
@@ -25,7 +22,7 @@ This directory contains comprehensive API design documentation for all core modu
 - Input preprocessing for various modalities
 - Output postprocessing and formatting
 
-### 4. [Executor Module API](executor_module_api.md)
+### Executor Module API
 **Executor module** coordinates and manages request execution across different engines and workers.
 
 **Key Components:**
@@ -34,7 +31,7 @@ This directory contains comprehensive API design documentation for all core modu
 - Worker management and task distribution
 - Error handling and recovery
 
-### 5. [Worker Module API](worker_module_api.md)
+### Worker Module API
 **Worker module** provides the actual execution environment for model inference.
 
 **Key Components:**
@@ -43,79 +40,6 @@ This directory contains comprehensive API design documentation for all core modu
 - Request batching and processing
 - Performance optimization
 
-## 🎯 How to Use These Templates
-
-### For New Module Development
-
-1. **Start with the Master Template**
-   - Copy `API_DESIGN_TEMPLATE.md`
-   - Rename to `[module_name]_api.md`
-   - Follow the structure exactly
-
-2. **Fill in Module-Specific Details**
-   - Update the module overview
-   - Define core classes and interfaces
-   - Specify public API methods
-   - Add configuration options
-   - Define error handling strategy
-   - Provide usage examples
-
-3. **Review and Validate**
-   - Ensure all sections are complete
-   - Verify code examples work
-   - Check for consistency with other modules
-   - Validate configuration options
-
-### For Existing Module Updates
-
-1. **Update API Documentation**
-   - Modify existing API files
-   - Add new methods and classes
-   - Update examples and configuration
-   - Maintain backward compatibility notes
-
-2. **Version Control**
-   - Track changes in git
-   - Use clear commit messages
-   - Tag major API changes
-
-## 📚 Template Structure
-
-Each API template follows this standardized structure:
-
-### 1. Module Overview
-- **Purpose**: What the module does
-- **Responsibilities**: Key responsibilities
-- **Dependencies**: Required modules
-- **Integration Points**: How it connects with other modules
-
-### 2. Core Classes/Interfaces
-- **Base Classes**: Abstract base classes
-- **Implementation Classes**: Concrete implementations
-- **Data Structures**: Key data models
-
-### 3. Public API Methods
-- **Initialization**: Constructor and setup
-- **Core Operations**: Main functionality
-- **Configuration**: Settings management
-- **Lifecycle Management**: Start/stop/cleanup
-- **Monitoring**: Status and metrics
-
-### 4. Configuration
-- **Configuration Classes**: Dataclasses for settings
-- **Required Parameters**: Must-have settings
-- **Optional Parameters**: Optional settings with defaults
-- **Validation**: Parameter validation rules
-
-### 5. Error Handling
-- **Custom Exceptions**: Module-specific errors
-- **Error Codes**: Standardized error codes
-- **Recovery Strategies**: Error recovery approaches
-
-### 6. Examples
-- **Basic Usage**: Simple usage examples
-- **Advanced Usage**: Complex scenarios
-- **Integration Examples**: Multi-module usage
 
 ## 🔧 Implementation Guidelines
 
@@ -155,11 +79,11 @@ Each API template follows this standardized structure:
 2. **Set Up Development Environment**
    ```bash
    # Clone the repository
-   git clone https://github.com/hsliuustc0106/vllm-omni.git
+   git clone https://github.com/vllm-project/vllm-omni.git
    cd vllm-omni
 
    # Install dependencies
-   pip install -r requirements-dev.txt
+   pip install -e ".[dev]"
 
    # Set up pre-commit hooks
    pre-commit install
@@ -189,28 +113,12 @@ Each API template follows this standardized structure:
    - Update documentation
    - Submit pull requests
 
-## 📖 Additional Resources
-
-### Architecture Overview
-- [vLLM-omni Architecture Design](../architecture/vLLM-omni%20arch%20design%20doc.md)
-- [Component Categorization](../../ARCHITECTURE_CATEGORIZATION.md)
-
-### Development Guides
-- [Development Setup](../../README.md#development)
-- [Testing Guidelines](../../tests/README.md)
-- [Contributing Guidelines](../../CONTRIBUTING.md)
-
-### Examples
-- [Basic Examples](../../examples/basic/)
-- [Advanced Examples](../../examples/advanced/)
-- [Multimodal Examples](../../examples/multimodal/)
-
 ## 🤝 Contributing
 
 We welcome contributions to improve the API documentation and implementation. Please:
 
 1. **Follow the Template Structure**
-   - Use the master template as a guide
+   - Use the [api template](https://github.com/vllm-project/vllm-omni/docs/contributing/design_documents/api_design_template.md) as a guide
    - Maintain consistency across modules
    - Include all required sections
 

docs/contributing/design_documents/architecture_overview.md

@@ -1,4 +1,4 @@
-# vLLM-omni Architecture Design Document
+# Architecture Overview
 
 # Introduction
 
@@ -114,7 +114,7 @@ vLLM-omni will import core modules from existing vLLM framework and introduce se
 
 For any questions or feedback regarding this design document, please contact  
 Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk),  
-Hongsheng Liu: [hsliuustc@gmail.com](mailto:hsliuustc@gmail.com)
+Hongsheng Liu: [vllm-project@gmail.com](mailto:vllm-project@gmail.com)
 
 [image1]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAALYAAAFwCAYAAAAVNoEZAABxw0lEQVR4XuydBVRcWb6ve90709OStnTScel4x93dIUaEuBJ3dzciRPDgEtzdCRpCBCIEEkIgQHCCBHe+t6uSvNvNnZ55r2/fng5zvrX+69TZp4qkan3nz2/vKup88sm/OWqXL3HxwnnOnj7JubOnuaB6jqtXLqJ28TxXRWmrX8Ha0gIzE0OuX1bFQFsTcxMDLMyMMDXS44bWVRxsLLG5aYKuxlVsLMxwc3bAWF8bNzdnHGwt8fb2wNXZHlcnW3HMDjsbC4Ju3cLO2hwvT3ecnexwsLcmIMAHHy83+X18PVxxtLMUt+1xcbTBw8URRwdr+f1tLM1wd3PB08NZlAs24ue4uTrKf0bD5yfxb8qWDWvZtnE9W0Vt37yBvTu3cWDPTlE7OHviCBfPnsTEwECIfg7Vk0e5fkkVoxtamBneQF1NlSsXTmMrZJaN6elcF7LexFaIZ6SngbOjvfy2l5DUSwhoa2mK9U1jfL09ibwTKaS3IPx2BA4ONnh7uuHr4yUX1d/Pi0B/X1xd7PHz8cTT3UUc85SfIMHBQUJ4a3EfH/mJ4OsjThoXB/njPNydJbEl3rFisTKrli5i9fLFrF21jC0bVdi2aR07tqzn+KF9nD91HF1NDU4eO8CZIwe5fO4UBjoacpEvnDnG5fMnsDQ1Rk/rOje0rwmxLeQd3chAC0d7W6zMjfBwc8JHiGlmpIu5oS6hISHcibgtOrE99+4+wE50dR8hraeHm1zOW4G+BAb4yjtyYICf/FiAED0wMICw0DDxODvR8QNEBeLv7yO6txPeotN7e7lLYku8Y8akccycOhGlGVNRVprJsoXzWLlEmbUrl7Jv5xbOHD8sBD/ApnWrOLZvN2ePH0FDRJVronMf3rdDiH1aiK4p9s+JE0AWS6xE59bA2EAHO6ub3DTRw1UI7Ovjg/oVEWV01Llz567ouH6i+wbz8FEMtiKaBPr7yzuvn6+nED9I3rH9/bwJDgqUSy3r1OHh4YSEBOPu7io6/W1CQ0PFSeAvxHeTd3B/X19JbIl3TB4zjOnjRzJz8jiUFCazYPZ0li1QYrXo4nt3bObE4QPs3LIRlRWLObh7h9jfj5rqGdGtj7N/11Yh9hk0r14W+yfQ0bgiZLZC8/pljPR1ROwwwdxYF2d7GyGfN1cunkJfiB0h4oe3lweRkfeIjn4iF9vf1wcnkbX9fL0ICb4ll1kutrgtkzdEyB4ZGUlQ0C35Y+/cuSPv3oHimOxkkHX2W6KjN3x+Ev+mjBzQmzGD+zJxxCAmjxrC9AmjUBZyL1dWEkJvEF17G6tEB1++aB47Nq0Xsm+Vy71PbHdv2yTP4LL9Q3t3onVdTYhsyCXV02LyqCuyuS762texlk0oXV05e/IQetrqBInuK5v8RUc/JiLiDvYiisgmf7Y2N0UssZJv1a9fFQJ7yiUPkXXrsBAePHiAn+jMsghy//590dlDRRTxk8cW2VhoWLAktsQ7Bv7UleH9egi5+8kFl3Xw+TOnsmTBLDGpXMsOMaFcunCuiChKbF27ml1CbtnEcvt6FbnY504eE517uxB+CxpXLwlxNbhwXjbhvCGXWFOMmZsZibztwMmj+7khsniA6Kwurs4iSoi87OKCuuj4urqa3DQ3xdHBVt6pXZwd0NJQJ1AmbGiw6NC3xYkQLbK6lxA6WC55yP8VWxZrbhEWHiKJLfGOPl06MKhnVwb06Cy2XZg0aqiIJJNYNHcG61cvl2dtZXF78byZrF22iHUrl7B53WqRwxcI6dfLO/WW9Wvk3f3i2VOonj7BeRFT9HW0ULtwRoh/FIMb2hgZ6otMLuvql9klos3COYri3xrOiEH96dujK907d+Snrp3o37c3Gzdu4MKFCxjqG7Bn9y7uRt6Vd+gHUQ9xFsKHiMmnbEy29ReRRbaaEhYWJo8qDZ+fxL8pfbt2pF+3TqJ+ZHCvLiKODGX2tAly8VRWLJFPJBcoKcpFX6E8j9WL300slyrPkXd0WdfeIE4A2UqKbHlQtpJy7tQxtNWvckZILYspOprX0dFSZ8+2DezcqMKsCaOZOGwQvTp1oHPbVnRo2YL2rZrz/TdN+O6rJrRp8YMQvA+q5y9w9NARbgUFEf3wEffuR2FrZyUXOiIiQnTyEPmkUSa2LNLcu/dAElviHb06taF357bIBB/cqzNjh/YXOXs0cxWnsWjebPlKyVwRTebNmILyLAWWzJ3NMuW5LJijIO/o61YvY/WShWxUWcmB3dvZLgQ/emAvZ04cZc/2zULmTaiKuLJv22ZmjxvJqH496dK6pagW9Gnfhp7tWtH5h+9p//23tPruG1qI+v6br/n2qy/47K9/odn337FGRQUHJ2fCbt/GyFhPxI4Q+eRRJrivr7d8AvngQZTo6tGS2BLv6P6jELtLO7nY/bv/yPCBPZkiBJylMIV5QuS5M6czR3EKStMnMU8m+5yZLFGazdwZU1mzfAmrly2UTyzXr1omjyMb1qxgn8jch/fvYdsGFXZsXMMusV0i7j+6V1eGdm3PGBF5pvXtztxBPZjetyujxL8/VPw/erRuzo9C8pZNheBNv+b7r7/iy88/o2PHjixcspTbtyO4IbK4bMnvzp13KySyzB0cIptYRktiS/wXnzdpzpdNmtGtQyt+6tiK/j06MXbEYCaPG8WMKROYOWUc0yaOY6ro4orittKMacxWmMysaeNZtmAei+fPEZ19FquWLmSdkHv5ovnyzL1dxJQ1Qvr1ohaOG8rxJTO4ukYJ/a1LODprDEsGdmdkux8Y2KopHb/+knZffk6LJp/xvRD5e3G7+Zdf0PKbr2gtJG8jZP/uqy9Zu24tqmdOEn77jrxjy96wkb1BExERyd2794m4c1cSW+Idn33ZlL99IaJA6x/o1KaZ6N7tGTmkP+NGDpXLPG38KLnkk8eNYPqkMXKpZ0ydIKQfK2LJHHkHV1aaIYReIM/jMtFlgss691KRzeeKn6G7YwVup3fgr3YAw83KLO/VkbX9OrNhYBfmd23JrC6tmNT+B7p89zVdRafu1uwbOjf9ivbffUUHsS/r4h3F/69Fs2YsWbhAZOn78owdGBgol1vWyWWTyYg7EZLYEu9o/cVXrO71E6aL53Nx1mT6iLw9sE9Phg/qx/jhQ+Q1doTYjhryXu6xTJNtx4scrjCVOSKiKImoIsvispJ19MULlEQGn8GsYf05vWI2+aGOZHqb4LJvNY+v7iP2+hEeXdpL2NF1uG9V5ti4nuwb1pUlPVoz88cWTGzblHGihrVuypA23zOsfXOGCPE7t/qBZt9+g7qmFmFh4XKxZXlb9kbPbdG1ZZ87afj8JP5NGd2mFfuHD8BCeQ43F8/hzMxxzB3al5H9fmJo/z4MH9CXEYMHMEpMKscNFxFl9AjGjxjERNHRFSbIJB/N9MljmDV9MjOnTUJh8nhmT5/CxNHDMNu/luI77pRGB5Dhb0t1zB2qnoRREuxEnpsx989sI3T/apzXKXBz6WRMl01Bf9F4VKcP5uj4Pmwd1p21g7ui3Lsj83p1oEerZrQV3fv7pt9y4uRxbgUHERIWRlh4BMFC9FvBoZLYEu/YObgfp0YO4ZLowtemjhViKXB1wVT2KIxBccRAxorOPVgIPnRQX8YMHSjv4GPeSz5p7EgmjhkhjywyoadPkmXxMfJMPmHoAOJstKhLjKLiaTiVCQ+pSY6jJj6K0jA3CjxNSdI/S4L2MR5c2sW9CzsIOqKC156lGC2bivr8MRye0I+dI3uyelBXlvXvxMA2zWnb7Ds+//QvjBkziiAxaQwJDRWTx3D5kqC/9Ja6xAcU27Rkcgvx67/59yzo2JrDo/thvVYZj51rCT20leCDm7DfvALN5XNRHj2UcQP7MKRPV9HJe4suLuvkQvZRw+SCTxg9XH57pDgZji6bSXVcJHWZidS+SaMuN43alGdysasfBlER6UNpsCNv/azJctQh1eoacdrHuX9hF67bFmK+SpFz04dwZHw/to3sxfqh3ZgssniXZl/LlwT/+p//ia+vn5g8+uPt7SX/YJSrq4sktsQ7un39Fd2/+ZoB337NpJZNWdezI/pKY3FSmU+o2mVCjc3xd/LCw9kHbSNrDl/QQFFpAaNHDGNov3dZfIzo3rLJ5tiR7/L4zJGDCBQ5ui4tntr8TOpF1WWnUJsUQ+2z+1RHBVF1L4DSUFeKAuzJdjEkzUaT+BtniLq8F+ftypitVOCsEPvoxH7sHt2XzUN7MLtHGwa0aEpLIfYXf/sMS0sLPDw8cHR0xN7ORl4Nn5/Evymd2rXgxzY/0EF0wj5dO3NE2xTj5CLMC+rQL6zDtKQe+/I6rEvrMS6owfBNFQY5leimlnLI0pOl2/aLDN6X0UPede9TG1dxS/sMiO5cHO5DcYgHed4OZNkZURbqSfWDYKrvBlN+25/iIDcK/Rx542JEpp0uTzWPEXl+Jw6b52O8fDrnFYZxcupg9o7pK7p2Dxb3aYdil5a0+eYrvvnicw4e3I/lTQsM9G+gr6eL3g1dSWyJd3T44Rt+bN2C+bsOc9w7ApMiuFkMN/JrsSyuwbm8FuuiWkzF/s3Caizf1mCZX4VpbgV6yW+5/iCRpVv3MkFMJBdMn0BehDf50SGQLaLH6xfUy+JHVCiVYZ6UeltRGexO9T2xHxFI0S2RtX0cyHMxI9P2Bo9Flw8/vQ3bjfMwWDqVc0LsE1MGsXd0H7aP6M6SPu2Z3bUVbb9pwjdffsnKlcvQ1rjO1SuXuHZVVpclsf8gNomSvdgfKuaXh//1qD3O4fLrKjRya9HJrUE3uxqr4mrcy2vwLKvDMKsKSyG0s9i3Lxad+20txqJj3xRlJeQ2yyhBNy6Dq0GPWLN8Kbn3g6jNEdEjL5uHphZcO3QRNyNHEnyCqY2NpfbFC2rjnlBxL0x0clkMuUmSmQZPdS9wR3U/t4Line truncated
 

docs/contributing/design_documents/index.md

@@ -0,0 +1,12 @@
+# Design Documents
+
+This section contains design documents and architecture specifications for vLLM-omni.
+
+## Architecture Documents
+
+- [vLLM-omni Software Design Document](vllm_omni_design.md)
+- [Architecture Overview](architecture_overview.md)
+
+## API Design Documentation
+
+- [vLLM-omni API Documentation](api_design_doc.md)

docs/contributing/design_documents/vllm_omni_design.md

@@ -1,4 +1,4 @@
-# vLLM-omni Software Design Document
+# Architecture Design Principles
 
 ## Overview
 
@@ -206,7 +206,7 @@ dit_stage_config.executor_class = DiffusersPipelineExecutor
 
 Similar to OmniLLM in offline inference, add some asynchronous processing, referring to AsyncLLM
 ```python
-from vllm.v1.engine.async_llm improt AsyncLLM
+from vllm.v1.engine.async_llm import AsyncLLM
 
 
 class AsyncOmniLLM(AsyncLLM):
@@ -317,13 +317,13 @@ class DiffusersPipelineExecutor(Executor):
 
         # Generate image
         prompt_embeds = self._get_and_process_prompt_embeds(scheduler_output, positive_magic)
-        negtive_prompt_embeds = self.pipe.embed_prompt(" ")
+        negative_prompt_embeds = self.pipe.embed_prompt(" ")
 
         width, height = aspect_ratios["16:9"]
 
         image = pipe(
             prompt_embeds=prompt_embeds,
-            negtive_prompt_embeds=negtive_prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
             width=width,
             height=height,
             num_inference_steps=50,

docs/index.md

@@ -0,0 +1,32 @@
+# Welcome to vLLM-omni
+
+**vLLM-omni** extends vLLM capabilities to support multi-modality model inference and serving, particularly focusing on non-autoregressive structures and non-textual outputs. It aims to support:
+
+- **Multi-modal Models**: Text, image, video, audio, and sensor data processing
+- **Non-autoregressive Architectures**: Diffusion Transformers (DiT) and other parallel generation models
+- **Heterogeneous Outputs**: Beyond traditional text generation to multimodal outputs
+
+## <span class="twemoji">🎯</span> Key Features
+
+### Multi-Engine Support
+
+- **Autoregressive Engine**: Traditional text generation with enhanced KV-caching
+- **Diffusion Engine**: Support for DiT models and iterative generation
+- **Hybrid Engine**: Combined AR+DiT processing pipelines
+
+### Modality Processing
+
+- **Text**: Advanced tokenization and embedding generation
+- **Image**: Vision encoder integration (CLIP, etc.)
+- **Audio**: Speech processing and audio embedding
+
+## <span class="twemoji">🤖</span> Supported Models
+
+- Qwen-omni (Thinker-Talker-Codec structure)
+- HunyunaImage 3.0 (Ongoing)
+- Qwen-Image (Ongoing)
+
+## <span class="twemoji">📚</span> Documentation Navigation
+
+- To run open-source models on vLLM-Omni, we recommend starting with the [:material-code-tags: User Quide](user_guide/getting_started/quickstart.md)
+- To develop and contribute to vLLM-Omni, we recommend starting with the [:material-tools: Developer Guide](contributing/README.md)

docs/mkdocs/hooks/generate_examples.py

@@ -0,0 +1,233 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-omni project
+import itertools
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+
+import regex as re
+
+logger = logging.getLogger("mkdocs")
+
+ROOT_DIR = Path(__file__).parent.parent.parent.parent
+ROOT_DIR_RELATIVE = "../../../../.."
+EXAMPLE_DIR = ROOT_DIR / "examples"
+EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
+
+
+def fix_case(text: str) -> str:
+    subs = {
+        "api": "API",
+        "cli": "CLI",
+        "cpu": "CPU",
+        "llm": "LLM",
+        "mae": "MAE",
+        "tpu": "TPU",
+        "gguf": "GGUF",
+        "lora": "LoRA",
+        "rlhf": "RLHF",
+        "vllm": "vLLM",
+        "omni": "Omni",
+        "openai": "OpenAI",
+        "lmcache": "LMCache",
+        "multilora": "MultiLoRA",
+        "mlpspeculator": "MLPSpeculator",
+        r"fp\d+": lambda x: x.group(0).upper(),  # e.g. fp16, fp32
+        r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
+    }
+    for pattern, repl in subs.items():
+        text = re.sub(rf"\b{pattern}\b", repl, text, flags=re.IGNORECASE)
+    return text
+
+
+@dataclass
+class Example:
+    """
+    Example class for generating documentation content from a given path.
+
+    Attributes:
+        path (Path): The path to the main directory or file.
+        category (str): The category of the document.
+        main_file (Path): The main file in the directory.
+        other_files (list[Path]): list of other files in the directory.
+        title (str): The title of the document.
+
+    Methods:
+        __post_init__(): Initializes the main_file, other_files, and title attributes.
+        determine_main_file() -> Path: Determines the main file in the given path.
+        determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
+        determine_title() -> str: Determines the title of the document.
+        generate() -> str: Generates the documentation content.
+    """  # noqa: E501
+
+    path: Path
+    category: str = None
+    main_file: Path = field(init=False)
+    other_files: list[Path] = field(init=False)
+    title: str = field(init=False)
+
+    def __post_init__(self):
+        self.main_file = self.determine_main_file()
+        self.other_files = self.determine_other_files()
+        self.title = self.determine_title()
+
+    @property
+    def is_code(self) -> bool:
+        return self.main_file.suffix != ".md"
+
+    def determine_main_file(self) -> Path:
+        """
+        Determines the main file in the given path.
+        If the path is a file, it returns the path itself. Otherwise, it searches
+        for Markdown files (*.md) in the directory and returns the first one found.
+        Returns:
+            Path: The main file path, either the original path if it's a file or the first
+            Markdown file found in the directory.
+        Raises:
+            IndexError: If no Markdown files are found in the directory.
+        """  # noqa: E501
+        return self.path if self.path.is_file() else list(self.path.glob("*.md")).pop()
+
+    def determine_other_files(self) -> list[Path]:
+        """
+        Determine other files in the directory excluding the main file.
+
+        This method checks if the given path is a file. If it is, it returns an empty list.
+        Otherwise, it recursively searches through the directory and returns a list of all
+        files that are not the main file.
+
+        Returns:
+            list[Path]: A list of Path objects representing the other files in the directory.
+        """  # noqa: E501
+        if self.path.is_file():
+            return []
+        # Text file extensions that can be safely included
+        text_extensions = {
+            ".py",
+            ".md",
+            ".sh",
+            ".yaml",
+            ".yml",
+            ".json",
+            ".txt",
+            ".toml",
+            ".cfg",
+            ".ini",
+        }
+        is_other_file = lambda file: (file.is_file() and file != self.main_file and file.suffix in text_extensions)  # noqa: E731
+        return [file for file in self.path.rglob("*") if is_other_file(file)]
+
+    def determine_title(self) -> str:
+        if not self.is_code:
+            # Specify encoding for building on Windows
+            with open(self.main_file, encoding="utf-8") as f:
+                first_line = f.readline().strip()
+            match = re.match(r"^#\s+(?P<title>.+)$", first_line)
+            if match:
+                return match.group("title")
+        return fix_case(self.path.stem.replace("_", " ").title())
+
+    def fix_relative_links(self, content: str) -> str:
+        """
+        Fix relative links in markdown content by converting them to gh-file
+        format.
+
+        Args:
+            content (str): The markdown content to process
+
+        Returns:
+            str: Content with relative links converted to gh-file format
+        """
+        # Regex to match markdown links [text](relative_path)
+        # This matches links that don't start with http, https, ftp, or #
+        link_pattern = r"\[([^\]]*)\]\((?!(?:https?|ftp)://|#)([^)]+)\)"
+
+        def replace_link(match):
+            link_text = match.group(1)
+            relative_path = match.group(2)
+
+            # Make relative to repo root
+            gh_file = (self.main_file.parent / relative_path).resolve()
+            gh_file = gh_file.relative_to(ROOT_DIR)
+
+            # Make GitHub URL
+            url = "https://github.com/vllm-project/vllm-omni/"
+            url += "tree/main" if self.path.is_dir() else "blob/main"
+            gh_url = f"{url}/{gh_file}"
+
+            return f"[{link_text}]({gh_url})"
+
+        return re.sub(link_pattern, replace_link, content)
+
+    def generate(self) -> str:
+        content = f"# {self.title}\n\n"
+        url = "https://github.com/vllm-project/vllm-omni/"
+        url += "tree/main" if self.path.is_dir() else "blob/main"
+        content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
+
+        # Use long code fence to avoid issues with
+        # included files containing code fences too
+        code_fence = "``````"
+
+        if self.is_code:
+            # Use relative path from ROOT_DIR for snippet inclusion
+            main_file_rel_path = self.main_file.relative_to(ROOT_DIR)
+            content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel_path}"\n{code_fence}\n'
+        else:
+            with open(self.main_file, encoding="utf-8") as f:
+                # Skip the title from md snippets as it's been included above
+                main_content = f.readlines()[1:]
+            content += self.fix_relative_links("".join(main_content))
+        content += "\n"
+
+        if not self.other_files:
+            return content
+
+        content += "## Example materials\n\n"
+        for file in sorted(self.other_files):
+            # Use relative path from ROOT_DIR for snippet inclusion
+            file_rel_path = file.relative_to(ROOT_DIR)
+            content += f'??? abstract "{file.relative_to(self.path)}"\n'
+            if file.suffix != ".md":
+                content += f"    {code_fence}{file.suffix[1:]}\n"
+            content += f'    --8<-- "{file_rel_path}"\n'
+            if file.suffix != ".md":
+                content += f"    {code_fence}\n"
+
+        return content
+
+
+def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
+    logger.info("Generating example documentation")
+    logger.debug("Root directory: %s", ROOT_DIR.resolve())
+    logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
+    logger.debug("Example document directory: %s", EXAMPLE_DOC_DIR.resolve())
+
+    # Create the EXAMPLE_DOC_DIR if it doesn't exist
+    if not EXAMPLE_DOC_DIR.exists():
+        EXAMPLE_DOC_DIR.mkdir(parents=True)
+
+    categories = sorted(p for p in EXAMPLE_DIR.iterdir() if p.is_dir())
+
+    examples = []
+    glob_patterns = ["*.py", "*.md"]
+    # Find categorised examples
+    for category in categories:
+        globs = [category.glob(pattern) for pattern in glob_patterns]
+        for path in itertools.chain(*globs):
+            examples.append(Example(path, category.stem))
+        # Find examples in subdirectories
+        for path in category.glob("*/*.md"):
+            examples.append(Example(path.parent, category.stem))
+
+    # Generate the example documentation
+    for example in sorted(examples, key=lambda e: e.path.stem):
+        example_name = f"{example.path.stem}.md"
+        doc_path = EXAMPLE_DOC_DIR / example.category / example_name
+        if not doc_path.parent.exists():
+            doc_path.parent.mkdir(parents=True)
+        # Specify encoding for building on Windows
+        with open(doc_path, "w+", encoding="utf-8") as f:
+            f.write(example.generate())
+        logger.debug("Example generated: %s", doc_path.relative_to(ROOT_DIR))

docs/mkdocs/stylesheets/extra.css

@@ -0,0 +1,55 @@
+/* Custom CSS for vLLM-omni documentation - based on vLLM's extra.css */
+
+/* Improve code block readability */
+.highlight {
+    border-radius: 4px;
+}
+
+/* Style API documentation */
+.md-typeset .doc-contents {
+    margin-bottom: 1.5em;
+}
+
+/* Improve table readability */
+.md-typeset table {
+    display: table;
+    width: 100%;
+}
+
+/* Light mode: darker section titles (from vLLM) */
+body[data-md-color-scheme="default"] .md-nav__item--section > label.md-nav__link .md-ellipsis {
+  color: rgba(0, 0, 0, 0.7) !important;
+  font-weight: 700;
+}
+
+/* Dark mode: lighter gray section titles (from vLLM) */
+body[data-md-color-scheme="slate"] .md-nav__item--section > label.md-nav__link .md-ellipsis {
+  color: rgba(255, 255, 255, 0.75) !important;
+  font-weight: 700;
+}
+
+/* Navigation: Expand first level of API Reference by default, collapse deeper levels */
+/* Show direct children of API Reference section (vllm_omni's children: config, core, entrypoints, etc.) */
+.md-nav__item--nested[data-md-state] > .md-nav {
+    display: none;
+}
+
+/* Expand active items */
+.md-nav__item--nested.md-nav__item--active > .md-nav {
+    display: block;
+}
+
+/* For API Reference: show first level children (config, core, entrypoints, etc.) by default */
+.md-nav__item--section:has([href*="api/vllm_omni/index"]) > .md-nav > .md-nav__list > .md-nav__item--nested > .md-nav {
+    display: block;
+}
+
+/* But keep deeper levels collapsed */
+.md-nav__item--section:has([href*="api/vllm_omni/index"]) > .md-nav > .md-nav__list > .md-nav__item--nested > .md-nav > .md-nav__list > .md-nav__item--nested > .md-nav {
+    display: none;
+}
+
+/* Expand deeper levels only when active */
+.md-nav__item--section:has([href*="api/vllm_omni/index"]) > .md-nav > .md-nav__list > .md-nav__item--nested > .md-nav > .md-nav__list > .md-nav__item--nested.md-nav__item--active > .md-nav {
+    display: block;
+}

docs/user_guide/examples/index.md

@@ -0,0 +1,8 @@
+# Examples
+
+vLLM-omni's examples are split into two categories:
+
+- If you are using vLLM-omni from within Python code, see the *Offline Inference* section.
+- If you are using vLLM-omni from an HTTP application or client, see the *Online Serving* section.
+
+For detailed example documentation, check the [examples directory](https://github.com/vllm-project/vllm-omni/tree/main/examples) in the repository.

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -0,0 +1,64 @@
+# Offline Example of vLLM-omni for Qwen2.5-omni
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples\offline_inference\qwen2_5_omni>.
+
+
+## 🛠️ Installation
+
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
+
+## Run examples (Qwen2.5-omni)
+### Multiple Prompts
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
+```bash
+tar -xf <Your Download Path>/seedtts_testset.tar
+cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
+python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
+  --input examples/offline_inference/qwen2_5_omni/meta.lst \
+  --output examples/offline_inference/qwen2_5_omni/top100.txt \
+  --topk 100
+```
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Then run the command below.
+```bash
+bash run_multiple_prompts.sh
+```
+### Single Prompts
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Then run the command below.
+```bash
+bash run_single_prompt.sh
+```
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples\offline_inference\qwen2_5_omni\end2end.py"
+    ``````
+??? abstract "extract_prompts.py"
+    ``````py
+    --8<-- "examples\offline_inference\qwen2_5_omni\extract_prompts.py"
+    ``````
+??? abstract "processing_omni.py"
+    ``````py
+    --8<-- "examples\offline_inference\qwen2_5_omni\processing_omni.py"
+    ``````
+??? abstract "run_multiple_prompts.sh"
+    ``````sh
+    --8<-- "examples\offline_inference\qwen2_5_omni\run_multiple_prompts.sh"
+    ``````
+??? abstract "run_single_prompt.sh"
+    ``````sh
+    --8<-- "examples\offline_inference\qwen2_5_omni\run_single_prompt.sh"
+    ``````
+??? abstract "utils.py"
+    ``````py
+    --8<-- "examples\offline_inference\qwen2_5_omni\utils.py"
+    ``````

docs/user_guide/examples/offline_inference/qwen_2_5_omni.md

@@ -1,5 +1,8 @@
 # Offline Example of vLLM-omni for Qwen2.5-omni
 
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen_2_5_omni>.
+
+
 ## 🛠️ Installation
 
 ### Environment setup
@@ -49,7 +52,7 @@ uv pip install --no-build-isolation --editable .
 ```
 
 ### Verification for successful installation of vLLM
-Just run the command below. If no error, it demonstrates that the installation is successfull.
+Just run the command below. If no error, it demonstrates that the installation is successful.
 ```bash
 python -c "import vllm._C"
 ```
@@ -59,34 +62,36 @@ Install additional requirements for vllm-omni
 ```bash
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm_omni
-uv pip install -r requirements.txt
+uv pip install -e .
 ```
 
 ## Run examples (Qwen2.5-omni)
-### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
-```bash
-tar -xf <Your Download Path>/seedtts_testset.tar
-cp seedtts_testset/en/meta.lst examples/offline_inference/qwen_2_5_omni/meta.lst
-python3 examples/offline_inference/qwen_2_5_omni/extract_prompts.py \
-  --input examples/offline_inference/qwen_2_5_omni/meta.lst \
-  --output examples/offline_inference/qwen_2_5_omni/top100.txt \
-  --topk 100
-```
+
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen_2_5_omni
 ```
-Modify PYTHONPATH in run_multiple_prompts.sh as your path of vllm_omni. Then run.
+Then run
 ```bash
-bash run_multiple_prompts.sh
-```
-### Single Prompts
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen_2_5_omni
-```
-Modify PYTHONPATH in run_single_prompt.sh as your path of vllm_omni. Then run.
-```bash
-bash run_single_prompt.sh
+bash run.sh
 ```
+The output audio is saved in ./output_audio
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen_2_5_omni/end2end.py"
+    ``````
+??? abstract "processing_omni.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen_2_5_omni/processing_omni.py"
+    ``````
+??? abstract "run.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/qwen_2_5_omni/run.sh"
+    ``````
+??? abstract "utils.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen_2_5_omni/utils.py"
+    ``````

docs/user_guide/examples/online_serving/README.md

@@ -0,0 +1,35 @@
+# Online serving Example of vLLM-omni for Qwen2.5-omni
+
+Source <https://github.com/vllm-project/vllm-omni/blob/main/examples\online_serving\README.md>.
+
+
+## 🛠️ Installation
+
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/blob/main/README.md)
+
+## Run examples (Qwen2.5-omni)
+
+Launch the server
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
+```
+
+If you have custom stage configs file, launch the server with command below
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+```
+
+Get into the example folder
+```bash
+cd examples/online_serving
+```
+
+Send request via python
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py
+```
+
+Send request via curl
+```bash
+bash run_curl_multimodal_generation.sh
+```

docs/user_guide/examples/online_serving/index.md

@@ -0,0 +1,5 @@
+# Online Serving
+
+## Available Examples
+
+Online serving is under active development. Coming soon.

docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md

@@ -0,0 +1,7 @@
+# OpenAI Chat Completion Client For Multimodal Generation
+
+Source <https://github.com/vllm-project/vllm-omni/blob/main/examples\online_serving\openai_chat_completion_client_for_multimodal_generation.py>.
+
+``````py
+--8<-- "examples\online_serving\openai_chat_completion_client_for_multimodal_generation.py"
+``````

docs/user_guide/getting_started/quickstart.md

@@ -0,0 +1,55 @@
+# Quickstart
+
+This guide will help you get started with vLLM-omni.
+
+## Installation
+
+Use Docker to keep consistent basic environment (Optional, Recommended)
+
+```bash
+docker run --gpus all --ipc=host --network=host -v $source_dir:$container_dir --rm --name $container_name -it nvcr.io/nvidia/pytorch:25.01-py3 bash
+```
+
+Set up basic uv environment
+
+```bash
+pip install uv
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+```
+
+Install vLLM
+
+```bash
+uv pip install vllm==0.11.0 --torch-backend=auto
+```
+
+Install vLLM-omni
+
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm_omni
+uv pip install -e .
+```
+
+## Run examples (Qwen2.5-omni)
+
+Get into the example folder
+
+```bash
+cd examples/offline_inference/qwen_2_5_omni
+```
+
+Then run.
+
+```bash
+bash run.sh
+```
+
+The output audio is saved in ./output_audio
+
+## Next Steps
+
+- Read the [architecture documentation](../../contributing/design_documents/vllm_omni_design.md)
+- Check out the [API reference](../../api/overview.md)
+- Explore the [examples](../examples/index.md)

examples/offline_inference/qwen2_5_omni/README.md

@@ -0,0 +1,34 @@
+# Offline Example of vLLM-omni for Qwen2.5-omni
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Run examples (Qwen2.5-omni)
+### Multiple Prompts
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
+```bash
+tar -xf <Your Download Path>/seedtts_testset.tar
+cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
+python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
+  --input examples/offline_inference/qwen2_5_omni/meta.lst \
+  --output examples/offline_inference/qwen2_5_omni/top100.txt \
+  --topk 100
+```
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Then run the command below.
+```bash
+bash run_multiple_prompts.sh
+```
+### Single Prompts
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Then run the command below.
+```bash
+bash run_single_prompt.sh
+```

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -45,29 +45,17 @@ def parse_args():
         default="Qwen/Qwen2.5-Omni-7B",
         help="Hugging Face repo id to download if needed.",
     )
-    parser.add_argument(
-        "--hf-revision", default=None, help="Optional HF revision (branch/tag/commit)."
-    )
-    parser.add_argument(
-        "--prompts", nargs="+", default=None, help="Input text prompts."
-    )
-    parser.add_argument(
-        "--voice-type", default="default", help="Voice type, e.g., m02, f030, default."
-    )
+    parser.add_argument("--hf-revision", default=None, help="Optional HF revision (branch/tag/commit).")
+    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
+    parser.add_argument("--voice-type", default="default", help="Voice type, e.g., m02, f030, default.")
     parser.add_argument(
         "--code2wav-dir",
         default=None,
         help="Path to code2wav folder (contains spk_dict.pt).",
     )
-    parser.add_argument(
-        "--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt)."
-    )
-    parser.add_argument(
-        "--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file."
-    )
-    parser.add_argument(
-        "--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"]
-    )
+    parser.add_argument("--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt).")
+    parser.add_argument("--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file.")
+    parser.add_argument("--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"])
     parser.add_argument("--max-model-len", type=int, default=32768)
     parser.add_argument(
         "--init-sleep-seconds",
@@ -109,10 +97,14 @@ def parse_args():
     parser.add_argument("--use-torchvision", action="store_true")
     parser.add_argument("--tokenize", action="store_true")
     parser.add_argument(
-        "--output-wav", default="output.wav", help="[Deprecated] Output wav directory (use --output-dir)."
+        "--output-wav",
+        default="output.wav",
+        help="[Deprecated] Output wav directory (use --output-dir).",
     )
     parser.add_argument(
-        "--output-dir", default="outputs", help="Output directory to save text and wav files together."
+        "--output-dir",
+        default="outputs",
+        help="Output directory to save text and wav files together.",
     )
     parser.add_argument(
         "--thinker-hidden-states-dir",
@@ -159,7 +151,7 @@ def main():
     try:
         # Preferred: load from txt file (one prompt per line)
         if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
-            with open(args.txt_prompts, "r", encoding="utf-8") as f:
+            with open(args.txt_prompts, encoding="utf-8") as f:
                 lines = [ln.strip() for ln in f.readlines()]
             args.prompts = [ln for ln in lines if ln != ""]
             print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
@@ -243,9 +235,7 @@ def main():
                 request_id = int(output.request_id)
                 audio_tensor = output.multimodal_output["audio"]
                 output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
-                sf.write(
-                    output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000
-                )
+                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
                 print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
 

examples/offline_inference/qwen2_5_omni/extract_prompts.py

@@ -1,30 +1,34 @@
 #!/usr/bin/env python3
 import argparse
 from typing import Optional
-import os
-import torch
 
 
 def extract_prompt(line: str) -> Optional[str]:
     # 提取第一个 '|' 与第二个 '|' 之间的内容
-    i = line.find('|')
+    i = line.find("|")
     if i == -1:
         return None
-    j = line.find('|', i + 1)
+    j = line.find("|", i + 1)
     if j == -1:
         return None
-    return line[i + 1:j].strip()
+    return line[i + 1 : j].strip()
 
 
 def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--input", "-i", required=True, help="Input .lst file path")
     parser.add_argument("--output", "-o", required=True, help="Output file path")
-    parser.add_argument("--topk", "-k", type=int, default=100, help="Extract the top K prompts (default: 100)")
+    parser.add_argument(
+        "--topk",
+        "-k",
+        type=int,
+        default=100,
+        help="Extract the top K prompts (default: 100)",
+    )
     args = parser.parse_args()
 
     prompts = []
-    with open(args.input, "r", encoding="utf-8", errors="ignore") as f:
+    with open(args.input, encoding="utf-8", errors="ignore") as f:
         for line in f:
             if len(prompts) >= args.topk:
                 break
@@ -39,5 +43,3 @@ def main() -> None:
 
 if __name__ == "__main__":
     main()
-
-

examples/offline_inference/qwen2_5_omni/processing_omni.py

@@ -72,8 +72,7 @@ def smart_resize(
     """
     if max(height, width) / min(height, width) > MAX_RATIO:
         raise ValueError(
-            f"absolute aspect ratio must be smaller than {MAX_RATIO}, "
-            f"got {max(height, width) / min(height, width)}"
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
         )
     h_bar = max(factor, round_by_factor(height, factor))
     w_bar = max(factor, round_by_factor(width, factor))
@@ -88,9 +87,7 @@ def smart_resize(
     return h_bar, w_bar
 
 
-def fetch_image(
-    ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR
-) -> Image.Image:
+def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
     if "image" in ele:
         image = ele["image"]
     else:
@@ -110,10 +107,7 @@ def fetch_image(
     else:
         image_obj = Image.open(image)
     if image_obj is None:
-        raise ValueError(
-            f"Unrecognized image input, support local path, http url, "
-            f"base64 and PIL.Image, got {image}"
-        )
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
     image = image_obj.convert("RGB")
     # resize
     if "resized_height" in ele and "resized_width" in ele:
@@ -163,9 +157,7 @@ def smart_nframes(
     Returns:
         int: the number of frames for video used for model inputs.
     """
-    assert not (
-        "fps" in ele and "nframes" in ele
-    ), "Only accept either `fps` or `nframes`"
+    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
     if "nframes" in ele:
         nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
     else:
@@ -179,10 +171,7 @@ def smart_nframes(
         nframes = min(max(nframes, min_frames), max_frames)
         nframes = round_by_factor(nframes, FRAME_FACTOR)
     if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
-        raise ValueError(
-            f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], "
-            f"but got {nframes}."
-        )
+        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
     return nframes
 
 
@@ -204,10 +193,7 @@ def _read_video_torchvision(
     video_path = ele["video"]
     if version.parse(torchvision.__version__) < version.parse("0.19.0"):
         if "http://" in video_path or "https://" in video_path:
-            warnings.warn(
-                "torchvision < 0.19.0 does not support http/https video path, "
-                "please upgrade to 0.19.0."
-            )
+            warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.")
         if "file://" in video_path:
             video_path = video_path[7:]
     st = time.time()
@@ -221,8 +207,7 @@ def _read_video_torchvision(
     total_frames, video_fps = video.size(0), info["video_fps"]
     total_duration = round(total_frames / video_fps, 3)
     logger.info(
-        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, "
-        f"duration={total_duration}s, time={time.time() - st:.3f}s"
+        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s"
     )
     nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
     idx = torch.linspace(0, total_frames - 1, nframes).round().long()
@@ -258,15 +243,10 @@ def _read_video_decord(
     vr = decord.VideoReader(video_path)
     # TODO: support start_pts and end_pts
     if "video_start" in ele or "video_end" in ele:
-        raise NotImplementedError(
-            "not support start_pts and end_pts in decord for now."
-        )
+        raise NotImplementedError("not support start_pts and end_pts in decord for now.")
     total_frames, video_fps = len(vr), vr.get_avg_fps()
     total_duration = round(total_frames / video_fps, 3)
-    logger.info(
-        f"decord:  {video_path=}, {total_frames=}, {video_fps=}, "
-        f"time={time.time() - st:.3f}s"
-    )
+    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
     nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
     idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
     video = vr.get_batch(idx).asnumpy()
@@ -295,9 +275,7 @@ def get_video_reader_backend() -> str:
     return video_reader_backend
 
 
-def fetch_video(
-    ele: dict, image_factor: int = IMAGE_FACTOR
-) -> torch.Tensor | list[Image.Image]:
+def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
     if isinstance(ele["video"], str):
         video_reader_backend = get_video_reader_backend()
         video, total_dur, nframes = VIDEO_READER_BACKENDS[video_reader_backend](ele)
@@ -341,9 +319,7 @@ def fetch_video(
         process_info.pop("type", None)
         process_info.pop("video", None)
         images = [
-            fetch_image(
-                {"image": video_element, **process_info}, size_factor=image_factor
-            )
+            fetch_image({"image": video_element, **process_info}, size_factor=image_factor)
             for video_element in ele["video"]
         ]
         nframes = ceil_by_factor(len(images), FRAME_FACTOR)

examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh

@@ -1,4 +1,3 @@
-export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH
 python end2end.py --model Qwen/Qwen2.5-Omni-7B \
                                  --voice-type "m02" \
                                  --dit-ckpt none \

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

@@ -1,4 +1,3 @@
-export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH
 python end2end.py --model Qwen/Qwen2.5-Omni-7B \
                                  --voice-type "m02" \
                                  --dit-ckpt none \

examples/offline_inference/qwen2_5_omni/utils.py

@@ -1,5 +1,5 @@
 import tempfile
-from typing import Dict, List, Optional, Union
+from typing import Optional, Union
 from urllib.request import urlopen
 
 import librosa
@@ -9,17 +9,17 @@ import soundfile as sf
 import torch
 import torchvision.io
 from processing_omni import fetch_image, fetch_video
-
+from transformers import AutoConfig, AutoProcessor
 from vllm.inputs import TextPrompt
+
 from vllm_omni.inputs.data import OmniTokensPrompt
 
 # Simple caches to avoid repeated heavy HF loads per prompt
-_PROCESSOR_CACHE: Dict[str, "AutoProcessor"] = {}
-_CONFIG_CACHE: Dict[str, "AutoConfig"] = {}
+_PROCESSOR_CACHE: dict[str, "AutoProcessor"] = {}
+_CONFIG_CACHE: dict[str, "AutoConfig"] = {}
 
 
 def get_system_prompt():
-
     return {
         "role": "system",
         "content": [
@@ -46,7 +46,6 @@ def resample_wav_to_16khz(input_filepath):
 
 
 def fetch_and_read_video(args, video_url: str, fps=2):
-
     def read_video_with_torchvision(video_file_name: str):
         video, audio, info = torchvision.io.read_video(
             video_file_name,
@@ -72,10 +71,8 @@ def fetch_and_read_video(args, video_url: str, fps=2):
         else:
             return video
 
-    def read_video_with_transformers(video_file_name: Union[str, List[str]]):
-        video, total_duration, nframes, second_per_grid = fetch_video(
-            {"video": video_file_name}
-        )
+    def read_video_with_transformers(video_file_name: Union[str, list[str]]):
+        video, total_duration, nframes, second_per_grid = fetch_video({"video": video_file_name})
         if total_duration is None and nframes is None:
             nframes = len(video)
             total_duration = 0.5 * nframes
@@ -95,8 +92,7 @@ def fetch_and_read_video(args, video_url: str, fps=2):
         with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
             resp = requests.get(video_url)
             assert resp.status_code == requests.codes.ok, (
-                f"Failed to fetch video from {video_url}, "
-                f"status_code:{resp.status_code}, resp:{resp}"
+                f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
             )
 
             temp_video_file.write(urlopen(video_url).read())
@@ -110,11 +106,10 @@ def fetch_and_read_video(args, video_url: str, fps=2):
 
 def make_inputs_qwen2_omni(
     args,
-    messages: List[Dict[str, Union[str, List[Dict[str, str]]]]],
+    messages: list[dict[str, Union[str, list[dict[str, str]]]]],
     use_audio_in_video: Optional[bool] = False,
     tokenize: bool = False,
 ) -> Union[OmniTokensPrompt, TextPrompt]:
-
     from transformers import AutoConfig, AutoProcessor
 
     # Cached processor/config to prevent per-prompt reloading and repeated warnings
@@ -133,7 +128,7 @@ def make_inputs_qwen2_omni(
     # Decide legacy flag only once based on config (default True if unknown)
     if getattr(args, "legacy_omni_video", None) is None:
         if config is not None and hasattr(config, "architectures"):
-            args.legacy_omni_video = not ("Qwen2_5OmniModel" in config.architectures)
+            args.legacy_omni_video = "Qwen2_5OmniModel" not in config.architectures
         else:
             args.legacy_omni_video = True
 
@@ -180,9 +175,7 @@ def make_inputs_qwen2_omni(
                     audios.append(resample_wav_to_16khz(ele[audio_key]))
                 else:
                     raise ValueError(f"Unknown ele {ele}")
-            elif use_audio_in_video and (
-                ele["type"] == "video" or ele["type"] == "video_url"
-            ):
+            elif use_audio_in_video and (ele["type"] == "video" or ele["type"] == "video_url"):
                 # use video as audio as well
                 if "video_url" in ele:
                     audio_key = "video_url"
@@ -198,7 +191,7 @@ def make_inputs_qwen2_omni(
                     audios.append(librosa.load(ele[audio_key], sr=16000)[0])
                     videos.append(fetch_and_read_video(args, audio_key))
                 else:
-                    raise ValueError("Unknown ele {}".format(ele))
+                    raise ValueError(f"Unknown ele {ele}")
                 # insert a audio after the video
                 message["content"].insert(
                     index + 1,
@@ -295,10 +288,7 @@ def make_audio_in_video_v2_prompt(args):
             "content": [
                 {
                     "type": "video_url",
-                    "video_url": (
-                        "https://qianwen-res.oss-cn-beijing.aliyuncs.com/"
-                        "Qwen2.5-Omni/draw_small.mp4"
-                    ),
+                    "video_url": ("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"),
                 },
             ],
         },
@@ -312,9 +302,7 @@ def make_audio_in_video_v2_prompt(args):
     return prompt
 
 
-def make_omni_prompt(
-    args, prompt=None
-) -> Union[OmniTokensPrompt, List[OmniTokensPrompt]]:
+def make_omni_prompt(args, prompt=None) -> Union[OmniTokensPrompt, list[OmniTokensPrompt]]:
     if args.prompt_type == "text":
         prompt = make_text_prompt(args, prompt)
     elif args.prompt_type == "audio-in-video-v2":

examples/online_serving/README.md

@@ -0,0 +1,32 @@
+# Online serving Example of vLLM-omni for Qwen2.5-omni
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../README.md)
+
+## Run examples (Qwen2.5-omni)
+
+Launch the server
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
+```
+
+If you have custom stage configs file, launch the server with command below
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+```
+
+Get into the example folder
+```bash
+cd examples/online_serving
+```
+
+Send request via python
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py
+```
+
+Send request via curl
+```bash
+bash run_curl_multimodal_generation.sh
+```

examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py

@@ -0,0 +1,103 @@
+import base64
+
+from openai import OpenAI
+
+# Modify OpenAI's API key and API base to use vLLM's API server.
+openai_api_key = "EMPTY"
+openai_api_base = "http://localhost:8091/v1"
+
+client = OpenAI(
+    # defaults to os.environ.get("OPENAI_API_KEY")
+    api_key=openai_api_key,
+    base_url=openai_api_base,
+)
+
+SEED = 42
+
+
+def get_system_prompt():
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def run_text_to_audio(model: str) -> None:
+    thinker_sampling_params = {
+        "temperature": 0.0,  # Deterministic - no randomness
+        "top_p": 1.0,  # Disable nucleus sampling
+        "top_k": -1,  # Disable top-k sampling
+        "max_tokens": 2048,
+        "seed": SEED,  # Fixed seed for sampling
+        "detokenize": True,
+        "repetition_penalty": 1.1,
+    }
+    talker_sampling_params = {
+        "temperature": 0.9,
+        "top_p": 0.8,
+        "top_k": 40,
+        "max_tokens": 2048,
+        "seed": SEED,  # Fixed seed for sampling
+        "detokenize": True,
+        "repetition_penalty": 1.05,
+        "stop_token_ids": [8294],
+    }
+    code2wav_sampling_params = {
+        "temperature": 0.0,  # Deterministic - no randomness
+        "top_p": 1.0,  # Disable nucleus sampling
+        "top_k": -1,  # Disable top-k sampling
+        "max_tokens": 2048,
+        "seed": SEED,  # Fixed seed for sampling
+        "detokenize": True,
+        "repetition_penalty": 1.1,
+    }
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        talker_sampling_params,
+        code2wav_sampling_params,
+    ]
+
+    chat_completion = client.chat.completions.create(
+        messages=[
+            get_system_prompt(),
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.",
+                    },
+                ],
+            },
+        ],
+        model=model,
+        extra_body={
+            "sampling_params_list": sampling_params_list
+        },  # Optional, it has a default setting in stage_configs of the corresponding model.
+    )
+
+    count = 0
+    for choice in chat_completion.choices:
+        if choice.message.audio:
+            audio_data = base64.b64decode(choice.message.audio.data)
+            audio_file_path = f"audio_{count}.wav"
+            with open(audio_file_path, "wb") as f:
+                f.write(audio_data)
+            print(f"Audio saved to {audio_file_path}")
+            count += 1
+        elif choice.message.content:
+            print("Chat completion output from text:", choice.message.content)
+
+
+if __name__ == "__main__":
+    run_text_to_audio("Qwen/Qwen2.5-Omni-7B")

examples/online_serving/run_curl_multimodal_generation.sh

@@ -0,0 +1,75 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SEED=42
+
+thinker_sampling_params='{
+  "temperature": 0.0,
+  "top_p": 1.0,
+  "top_k": -1,
+  "max_tokens": 2048,
+  "seed": 42,
+  "detokenize": true,
+  "repetition_penalty": 1.1
+}'
+
+talker_sampling_params='{
+  "temperature": 0.9,
+  "top_p": 0.8,
+  "top_k": 40,
+  "max_tokens": 2048,
+  "seed": 42,
+  "detokenize": true,
+  "repetition_penalty": 1.05,
+  "stop_token_ids": [8294]
+}'
+
+code2wav_sampling_params='{
+  "temperature": 0.0,
+  "top_p": 1.0,
+  "top_k": -1,
+  "max_tokens": 2048,
+  "seed": 42,
+  "detokenize": true,
+  "repetition_penalty": 1.1
+}'
+# Above is optional, it has a default setting in stage_configs of the corresponding model.
+
+output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @- <<EOF
+{
+  "model": "Qwen/Qwen2.5-Omni-7B",
+  "extra_body": {
+    "sampling_params_list": [
+      $thinker_sampling_params,
+      $talker_sampling_params,
+      $code2wav_sampling_params
+    ]
+  },
+  "messages": [
+    {
+      "role": "system",
+      "content": [
+        {
+          "type": "text",
+          "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
+        }
+      ]
+    },
+    {
+      "role": "user",
+      "content": [
+        {
+          "type": "text",
+          "text": "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+        }
+      ]
+    }
+  ]
+}
+EOF
+)
+
+# Here it only shows the text content of the first choice. Audio content has many binaries, so it's not displayed here.
+echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"

mkdocs.yml

@@ -0,0 +1,180 @@
+site_name: vLLM-omni
+site_description: Multi-modality models inference and serving with non-autoregressive structures
+site_author: vLLM-omni Team
+site_url: https://vllm-project.github.io/vllm-omni/
+
+repo_name: vllm-project/vllm-omni
+repo_url: https://github.com/vllm-project/vllm-omni
+edit_uri: edit/main/docs/
+
+# Copyright
+copyright: Copyright © 2025 vLLM-omni Team
+
+# Theme
+theme:
+  name: material
+  palette:
+    # Light mode
+    - scheme: default
+      primary: blue
+      accent: blue
+      toggle:
+        icon: material/brightness-7
+        name: Switch to dark mode
+    # Dark mode
+    - scheme: slate
+      primary: blue
+      accent: blue
+      toggle:
+        icon: material/brightness-4
+        name: Switch to light mode
+  features:
+    - content.action.edit
+    - content.code.copy
+    - content.tabs.link
+    - navigation.instant
+    - navigation.instant.progress
+    - navigation.tracking
+    - navigation.tabs
+    - navigation.tabs.sticky
+    - navigation.sections
+    - navigation.indexes
+    - navigation.top
+    - search.suggest
+    - search.highlight
+    - search.share
+    - content.code.annotate
+    - content.tabs
+    - content.tooltips
+    - toc.follow
+  icon:
+    repo: fontawesome/brands/github
+
+hooks:
+  - docs/mkdocs/hooks/generate_examples.py
+
+# Plugins
+plugins:
+  - meta
+  - search
+  # - awesome-nav  # Uncomment if installed
+  - glightbox
+  - git-revision-date-localized:
+      # exclude files
+      exclude:
+        - api/*
+        - user_guide/examples/**
+        - contributing/design_documents/api_design_template.md
+        - DOCS_GUIDE.md
+  - minify:
+      minify_html: true
+      minify_js: true
+      minify_css: true
+      cache_safe: true
+  - api-autonav:
+      modules: ["vllm_omni"]
+      api_root_uri: "api"
+      nav_item_prefix: ""  # No prefix in navigation tree (clean names)
+      show_full_namespace: false  # Show only module name, not full path
+      exclude:
+        - "re:vllm_omni\\._.*"  # Internal modules
+  - mkdocstrings:
+      handlers:
+        python:
+          paths: [vllm_omni]
+          options:
+            show_symbol_type_heading: true
+            show_symbol_type_toc: true
+            show_source: true
+            show_root_heading: true
+            show_root_toc_entry: true
+            show_root_full_path: false
+            show_object_full_path: false
+            show_if_no_docstring: true
+            show_signature_annotations: true
+            separate_signature: true
+            show_overloads: true
+            signature_crossrefs: true
+            filters: []
+            summary:
+              modules: true
+          inventories:
+            - https://docs.python.org/3/objects.inv
+            - https://typing-extensions.readthedocs.io/en/latest/objects.inv
+            - https://docs.aiohttp.org/en/stable/objects.inv
+            - https://pillow.readthedocs.io/en/stable/objects.inv
+            - https://numpy.org/doc/stable/objects.inv
+            - https://pytorch.org/docs/stable/objects.inv
+            - https://psutil.readthedocs.io/en/stable/objects.inv
+
+# Markdown extensions
+markdown_extensions:
+  - attr_list
+  - md_in_html
+  - admonition
+  - pymdownx.details
+  # For content tabs
+  - pymdownx.superfences
+  - pymdownx.tabbed:
+      slugify: !!python/object/apply:pymdownx.slugs.slugify
+        kwds:
+          case: lower
+      alternate_style: true
+  # For code highlighting
+  - pymdownx.highlight:
+      anchor_linenums: true
+      line_spans: __span
+      pygments_lang_class: true
+  - pymdownx.inlinehilite
+  - pymdownx.snippets
+  # For emoji and icons
+  - pymdownx.emoji:
+      emoji_index: !!python/name:material.extensions.emoji.twemoji
+      emoji_generator: !!python/name:material.extensions.emoji.to_svg
+  # For in page [TOC] (not sidebar)
+  - toc:
+      permalink: true
+  # For math rendering
+  - pymdownx.arithmatex:
+      generic: true
+  # For task lists
+  - pymdownx.tasklist:
+      custom_checkbox: true
+  # For tables
+  - tables
+
+# Navigation
+nav:
+  - Home: index.md
+  - User Guide:
+    - Getting Started:
+      - user_guide/getting_started/quickstart.md
+    - Examples:
+      - user_guide/examples/index.md
+      - Offline Inference:
+        - user_guide/examples/offline_inference/qwen_2_5_omni.md
+      - Online Serving:
+        - user_guide/examples/online_serving/index.md
+  - API Reference:
+      - Summary: api/overview.md
+  - Developer Guide:
+    - contributing/README.md
+    - Design Documents:
+      - contributing/design_documents/index.md
+      - contributing/design_documents/vllm_omni_design.md
+      - contributing/design_documents/architecture_overview.md
+      - contributing/design_documents/api_design_doc.md
+
+# Extra
+extra:
+  version:
+    provider: mike
+
+# Extra CSS
+extra_css:
+  - mkdocs/stylesheets/extra.css
+
+# Makes the url format end in .html rather than act as a dir
+# So index.md generates as index.html and is available under URL /index.html
+# https://www.mkdocs.org/user-guide/configuration/#use_directory_urls
+use_directory_urls: false

pyproject.toml

@@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vllm-omni"
-version = "0.1.0"
+version = "0.11.0rc1"
 description = "vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures"
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.9,<3.14"  # Align with vLLM v0.11 supported Python versions
 license = {text = "Apache-2.0"}
 authors = [
-    {name = "vLLM-omni Team", email = "hsliuustc@gmail.com"}
+    {name = "vLLM-omni Team", email = "vllm-project@gmail.com"}
 ]
 keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
 classifiers = [
@@ -19,57 +19,55 @@ classifiers = [
     "Intended Audience :: Science/Research",
     "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 
 dependencies = [
-    "vllm>=0.2.0",
-    "torch>=2.0.0",
-    "transformers>=4.30.0",
-    "numpy>=1.21.0",
-    "pillow>=8.0.0",
-    "opencv-python>=4.5.0",
-    "PyYAML>=6.0",
-    "gradio>=3.0.0",
-    "fastapi>=0.100.0",
-    "uvicorn>=0.20.0",
-    "pydantic>=2.0.0",
-    "asyncio-mqtt>=0.11.0",
-    "ray>=2.0.0",
+    # Core runtime dependencies (required for actual usage)
+    "omegaconf>=2.3.0",
+    "decord>=0.6.0",
+    "librosa>=0.11.0",
+    "resampy>=0.4.3",
+    # "vllm==0.11.0", # install manually
 ]
 
 [project.optional-dependencies]
+
 dev = [
     "pytest>=7.0.0",
     "pytest-asyncio>=0.21.0",
     "pytest-cov>=4.0.0",
     "black>=23.0.0",
     "isort>=5.12.0",
-    "flake8>=6.0.0",
     "mypy>=1.0.0",
     "pre-commit>=3.0.0",
 ]
+
 docs = [
-    "sphinx>=6.0.0",
-    "sphinx-rtd-theme>=1.2.0",
-    "myst-parser>=1.0.0",
-]
-examples = [
-    "jupyter>=1.0.0",
-    "matplotlib>=3.5.0",
-    "seaborn>=0.11.0",
+    "mkdocs>=1.5.0",
+    "mkdocs-api-autonav",
+    "mkdocs-material",
+    "mkdocstrings-python",
+    "mkdocs-gen-files",
+    # "mkdocs-awesome-nav",  # Not available on PyPI
+    "mkdocs-glightbox",
+    "mkdocs-git-revision-date-localized-plugin",
+    "mkdocs-minify-plugin",
+    "regex",
+    "pydantic",
 ]
 
+
 [project.urls]
-Homepage = "https://github.com/hsliuustc0106/vllm-omni"
-Repository = "https://github.com/hsliuustc0106/vllm-omni"
+Homepage = "https://github.com/vllm-project/vllm-omni"
+Repository = "https://github.com/vllm-project/vllm-omni"
 Documentation = "https://vllm-omni.readthedocs.io"
-"Bug Tracker" = "https://github.com/hsliuustc0106/vllm-omni/issues"
+"Bug Tracker" = "https://github.com/vllm-project/vllm-omni/issues"
 
 [project.scripts]
 vllm = "vllm_omni.entrypoints.cli.main:main"
@@ -79,36 +77,44 @@ vllm-omni = "vllm_omni.entrypoints.cli.main:main"
 where = ["."]
 include = ["vllm_omni*"]
 
-[tool.black]
-line-length = 88
-target-version = ['py38', 'py39', 'py310', 'py311']
-include = '\.pyi?$'
-extend-exclude = '''
-/(
-  # directories
-  \.eggs
-  | \.git
-  | \.hg
-  | \.mypy_cache
-  | \.tox
-  | \.venv
-  | build
-  | dist
-)/
-'''
+[tool.ruff]
+line-length = 120
+exclude = [
+    ".eggs",
+    ".git",
+    ".hg",
+    ".mypy_cache",
+    ".tox",
+    ".venv",
+    "build",
+    "dist",
+    "vllm_omni.egg-info",
+]
 
-[tool.isort]
-profile = "black"
-multi_line_output = 3
-line_length = 88
-known_first_party = ["vllm_omni"]
+[tool.ruff.lint]
+select = [
+    "E",  # pycodestyle errors
+    "W",  # pycodestyle warnings
+    "F",  # pyflakes
+    "I",  # isort (handled separately, but included for compatibility)
+    "N",  # pep8-naming
+    "UP", # pyupgrade
+]
+ignore = [
+    "E203",  # whitespace before ':' (conflicts with black)
+    # W503 is not needed in ruff as it's compatible with black by default
+    "N801",  # class names should use CapWords convention
+    "N802",  # function name should follow snake_case
+    "N806",  # variable in function should follow snake_case
+    "N812",  # lowercase imported as non-lowercase: functional as F
+]
 
-[tool.flake8]
-max-line-length = 88
-extend-ignore = ["E203", "W503"]
+[tool.ruff.lint.per-file-ignores]
+"examples/**" = ["E501"]  # Allow long lines in examples
+"tests/**" = ["E501"]  # Allow long lines in tests
 
 [tool.mypy]
-python_version = "3.8"
+python_version = "3.12, 3.13"
 warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true
@@ -141,3 +147,11 @@ markers = [
     "benchmark: Benchmark tests",
     "slow: Slow tests",
 ]
+
+[tool.typos.default]
+extend-ignore-identifiers-re = [
+    ".*_thw",
+    ".*thw",
+    "ein",
+    ".*arange"
+]

requirements.txt

@@ -1,4 +0,0 @@
-omegaconf==2.3.0
-decord==0.6.0
-librosa==0.11.0
-resampy==0.4.3

vllm_omni/__init__.py

@@ -12,16 +12,28 @@ Architecture:
   processing
 """
 
-from . import patch  # noqa: F401
+try:
+    from . import patch  # noqa: F401
+except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
+    if exc.name != "vllm":
+        raise
+    # Allow importing vllm_omni without vllm (e.g., documentation builds)
+    patch = None  # type: ignore
 from .config import OmniModelConfig
+from .entrypoints.async_omni_llm import AsyncOmniLLM
 
 # Main entry points
 from .entrypoints.omni_llm import OmniLLM
 
+# from .version import __version__, __version_tuple__  # isort:skip
+
+
 __all__ = [
+    # "__version__",  # TODO: add __version__ back when it is implemented
+    # "__version_tuple__",  # TODO: add __version_tuple__ back when it is implemented
     # Main components
     "OmniLLM",
-    # "AsyncOmniLLM", # TODO: add AsyncOmniLLM back when it is implemented
+    "AsyncOmniLLM",
     # Configuration
     "OmniModelConfig",
     # All other components are available through their respective modules

vllm_omni/config/model.py

@@ -2,9 +2,9 @@ from typing import Optional
 
 from pydantic import ConfigDict
 from pydantic.dataclasses import dataclass
+from vllm.config import ModelConfig, config
 
 import vllm_omni.model_executor.models as me_models
-from vllm.config import ModelConfig, config
 
 
 @config
@@ -15,6 +15,7 @@ class OmniModelConfig(ModelConfig):
     stage_id: int = 0
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
+    ode_solver_class: Optional[str] = None
     engine_output_type: Optional[str] = None
 
     @property

vllm_omni/core/sched/diffusion_scheduler.py

@@ -4,15 +4,10 @@ from typing import Optional
 
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.request_queue import create_request_queue
-from vllm.v1.core.sched.scheduler import (
-    EngineCoreOutputs,
-    Request,
-    RequestStatus,
-    SchedulerOutput,
-    SpecDecodingStats,
-)
+from vllm.v1.core.sched.scheduler import EngineCoreOutputs, Request, RequestStatus, SchedulerOutput, SpecDecodingStats
 from vllm.v1.core.sched.utils import remove_all
 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
+
 from vllm_omni.core.sched.output import OmniNewRequestData
 from vllm_omni.core.sched.scheduler import OmniScheduler
 from vllm_omni.outputs import OmniModelRunnerOutput
@@ -97,17 +92,13 @@ class DiffusionScheduler(OmniScheduler):
         num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
         if self.running:
             any_request = self.running[0]
-            num_common_prefix_blocks = (
-                self.kv_cache_manager.get_num_common_prefix_blocks(
-                    any_request, len(self.running)
-                )
+            num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(
+                any_request, len(self.running)
             )
 
         # Assemble SchedulerOutput
         new_reqs_data = [
-            OmniNewRequestData.from_request(
-                req, req_to_new_blocks[req.request_id].get_block_ids()
-            )
+            OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids())
             for req in scheduled_new_reqs
         ]
         cached_reqs_data = self._make_cached_request_data(
@@ -117,9 +108,7 @@ class DiffusionScheduler(OmniScheduler):
             scheduled_spec_decode_tokens,
             req_to_new_blocks,
         )
-        scheduled_requests = (
-            scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
-        )
+        scheduled_requests = scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
         structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(
             scheduled_requests, scheduled_spec_decode_tokens
         )
@@ -177,9 +166,7 @@ class DiffusionScheduler(OmniScheduler):
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: Optional[SpecDecodingStats] = None
-        kv_connector_stats = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
+        kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
 
         # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
         # the below loop can be a performance bottleneck. We should do our best
@@ -196,13 +183,9 @@ class DiffusionScheduler(OmniScheduler):
                 continue
 
             req_index = model_runner_output.req_id_to_index[req_id]
-            generated_token_ids = (
-                sampled_token_ids[req_index] if sampled_token_ids else []
-            )
+            generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
 
-            scheduled_spec_token_ids = (
-                scheduler_output.scheduled_spec_decode_tokens.get(req_id)
-            )
+            scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
             if scheduled_spec_token_ids:
                 num_draft_tokens = len(scheduled_spec_token_ids)
                 num_accepted = len(generated_token_ids) - 1
@@ -239,11 +222,7 @@ class DiffusionScheduler(OmniScheduler):
                 stopped_preempted_reqs.add(request)
 
             # Extract sample logprobs if needed.
-            if (
-                request.sampling_params is not None
-                and request.sampling_params.logprobs is not None
-                and logprobs
-            ):
+            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
                 # NOTE: once we support N tokens per step (spec decode),
                 # the outer lists can be of length > 1.
                 new_logprobs = logprobs.slice(req_index, req_index + 1)
@@ -252,9 +231,7 @@ class DiffusionScheduler(OmniScheduler):
                 # NOTE: structured_output_request should not be None if
                 # use_structured_output, we have check above, so safe to ignore
                 # type warning
-                request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]  # noqa: E501
-                    req_id, new_token_ids
-                )
+                request.structured_output_request.grammar.accept_tokens(req_id, new_token_ids)  # type: ignore[union-attr]  # noqa: E501
 
             # spec_token_ids comes from the model runner output
             if num_nans_in_logits is not None and req_id in num_nans_in_logits:
@@ -263,7 +240,6 @@ class DiffusionScheduler(OmniScheduler):
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
             if new_token_ids or pooler_output is not None or kv_transfer_params:
-
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
                     EngineCoreOutput(
@@ -298,10 +274,7 @@ class DiffusionScheduler(OmniScheduler):
 
         # Create EngineCoreOutputs for all clients that have requests with
         # outputs in this step.
-        engine_core_outputs = {
-            client_index: EngineCoreOutputs(outputs=outs)
-            for client_index, outs in outputs.items()
-        }
+        engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
 
         finished_req_ids = self.finished_req_ids_dict
         if finished_req_ids:
@@ -312,14 +285,10 @@ class DiffusionScheduler(OmniScheduler):
                 if (eco := engine_core_outputs.get(client_index)) is not None:
                     eco.finished_requests = finished_set
                 else:
-                    engine_core_outputs[client_index] = EngineCoreOutputs(
-                        finished_requests=finished_set
-                    )
+                    engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
             finished_req_ids.clear()
 
-        if (
-            stats := self.make_stats(spec_decoding_stats, kv_connector_stats)
-        ) is not None:
+        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
             # Return stats to only one of the front-ends.
             if (eco := next(iter(engine_core_outputs.values()), None)) is None:
                 # We must return the stats even if there are no request

vllm_omni/core/sched/output.py

@@ -3,6 +3,7 @@ from typing import Optional
 
 from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.request import Request
+
 from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
 
 

vllm_omni/core/sched/scheduler.py

@@ -39,22 +39,14 @@ class OmniScheduler(VLLMScheduler):
                     num_computed_tokens=nr.num_computed_tokens,
                     lora_request=nr.lora_request,
                     # Enrich with omni payloads from the live request object
-                    prompt_embeds=(
-                        getattr(request, "prompt_embeds", None) if request else None
-                    ),
-                    additional_information=(
-                        getattr(request, "additional_information", None)
-                        if request
-                        else None
-                    ),
+                    prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
+                    additional_information=(getattr(request, "additional_information", None) if request else None),
                 )
                 new_list.append(omni_nr)
 
             scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
         except Exception:
             # If anything goes wrong, leave the original output unchanged
-            init_logger(__name__).exception(
-                "Failed to wrap scheduled_new_reqs with OmniNewRequestData"
-            )
+            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
 
         return scheduler_output

vllm_omni/engine/__init__.py

@@ -5,7 +5,6 @@ Engine components for vLLM-omni.
 from typing import Any, Optional
 
 import msgspec
-
 from vllm.v1.engine import EngineCoreRequest
 
 

vllm_omni/engine/arg_utils.py

@@ -2,11 +2,43 @@ from dataclasses import dataclass
 from typing import Optional
 
 from vllm.engine.arg_utils import EngineArgs
+from vllm.v1.engine.async_llm import AsyncEngineArgs
+
 from vllm_omni.config import OmniModelConfig
 
 
 @dataclass
 class OmniEngineArgs(EngineArgs):
+    stage_id: int = 0
+    model_stage: str = "thinker"
+    model_arch: str = "Qwen2_5OmniForConditionalGeneration"
+    ode_solver_class: Optional[str] = None
+    engine_output_type: Optional[str] = None
+
+    def create_model_config(self) -> OmniModelConfig:
+        # First, get the base ModelConfig from the parent class
+        base_config = super().create_model_config()
+
+        # Create OmniModelConfig by copying all base config attributes
+        # and adding the new omni-specific fields
+        config_dict = base_config.__dict__.copy()
+
+        # Add the new omni-specific fields
+        config_dict["stage_id"] = self.stage_id
+        config_dict["model_stage"] = self.model_stage
+        config_dict["model_arch"] = self.model_arch
+        config_dict["engine_output_type"] = self.engine_output_type
+        config_dict["ode_solver_class"] = self.ode_solver_class
+
+        # Create and return the OmniModelConfig instance
+        omni_config = OmniModelConfig(**config_dict)
+        omni_config.hf_config.architectures = omni_config.architectures
+
+        return omni_config
+
+
+@dataclass
+class AsyncOmniEngineArgs(AsyncEngineArgs):
     stage_id: int = 0
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"

vllm_omni/engine/output_processor.py

@@ -1,7 +1,6 @@
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import torch
-
 from vllm.logger import init_logger
 from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
@@ -10,20 +9,16 @@ from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
 from vllm.v1.engine.detokenizer import IncrementalDetokenizer
 from vllm.v1.engine.logprobs import LogprobsProcessor
 from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
-from vllm.v1.engine.output_processor import (
-    OutputProcessorOutput,
-    RequestOutputCollector,
-    RequestState,
-)
+from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestOutputCollector, RequestState
 from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.metrics.stats import IterationStats
+
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
 
 
 class OmniRequestState(RequestState):
-
     def __init__(
         self,
         *args,
@@ -44,7 +39,6 @@ class OmniRequestState(RequestState):
         queue: Optional[Any],
         log_stats: bool,
     ) -> "OmniRequestState":
-
         if sampling_params := request.sampling_params:
             if not sampling_params.detokenize:
                 tokenizer = None
@@ -75,9 +69,7 @@ class OmniRequestState(RequestState):
             request_id=request.request_id,
             parent_req=parent_req,
             request_index=request_index,
-            lora_name=(
-                request.lora_request.name if request.lora_request is not None else None
-            ),
+            lora_name=(request.lora_request.name if request.lora_request is not None else None),
             output_kind=output_kind,
             prompt=prompt,
             prompt_token_ids=request.prompt_token_ids,
@@ -93,9 +85,7 @@ class OmniRequestState(RequestState):
             log_stats=log_stats,
         )
 
-    def add_multimodal_tensor(
-        self, tensor: Optional[torch.Tensor], mm_type: Optional[str]
-    ) -> None:
+    def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], mm_type: Optional[str]) -> None:
         if tensor is None:
             return
         try:
@@ -137,15 +127,11 @@ class OmniRequestState(RequestState):
         if self.parent_req is None:
             outputs = [output]
         else:
-            request_id, outputs, finished = self.parent_req.get_outputs(
-                request_id, output
-            )
+            request_id, outputs, finished = self.parent_req.get_outputs(request_id, output)
             if not outputs:
                 return None
 
-        return self._new_request_output(
-            request_id, outputs, finished, kv_transfer_params
-        )
+        return self._new_request_output(request_id, outputs, finished, kv_transfer_params)
 
     def _new_completion_output(
         self,
@@ -154,9 +140,7 @@ class OmniRequestState(RequestState):
         stop_reason: Optional[Union[int, str]],
     ) -> Any:
         # Reuse base text/logprobs logic, then annotate with pooling_result.
-        base_output = super()._new_completion_output(
-            token_ids, finish_reason, stop_reason
-        )
+        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
         try:
             if self.mm_accumulated is not None:
                 tensor = self.mm_accumulated
@@ -196,14 +180,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         engine_core_output_type: Optional[str] = None,
     ):
         super().__init__(tokenizer=tokenizer, log_stats=log_stats)
-        self.output_handlers: Dict[str, Callable[[EngineCoreOutput], None]] = {}
-        self._reqid_to_mm_type: Dict[str, str] = {}
+        self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
+        self._reqid_to_mm_type: dict[str, str] = {}
         self.request_states: dict[str, OmniRequestState] = {}
         self.engine_core_output_type = engine_core_output_type
 
-    def register_handler(
-        self, modality: str, handler: Callable[[EngineCoreOutput], None]
-    ) -> None:
+    def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None:
         self.output_handlers[modality.lower()] = handler
 
     def add_request(
@@ -255,9 +237,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 continue
 
             # 1) Stats
-            self._update_stats_from_output(
-                req_state, eco, engine_core_timestamp, iteration_stats
-            )
+            self._update_stats_from_output(req_state, eco, engine_core_timestamp, iteration_stats)
 
             new_token_ids = eco.new_token_ids
             pooling_output = eco.pooling_output
@@ -270,9 +250,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             # 2) Detokenize and logprobs when text path
             assert req_state.detokenizer is not None
             assert req_state.logprobs_processor is not None
-            stop_string = req_state.detokenizer.update(
-                new_token_ids, finish_reason == FinishReason.STOP
-            )
+            stop_string = req_state.detokenizer.update(new_token_ids, finish_reason == FinishReason.STOP)
             if stop_string:
                 finish_reason = FinishReason.STOP
                 stop_reason = stop_string
@@ -280,12 +258,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
 
             # 2.5) Accumulate multimodal tensors in RequestState
             try:
-                mm_type = (
-                    getattr(eco, "output_type", self.engine_core_output_type) or ""
-                ).lower()
-                if pooling_output is not None and isinstance(
-                    req_state, OmniRequestState
-                ):
+                mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
+                if pooling_output is not None and isinstance(req_state, OmniRequestState):
                     req_state.add_multimodal_tensor(pooling_output, mm_type)
             except Exception:
                 logger.debug(
@@ -310,18 +284,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             if ro:
                 # Attach accumulated multimodal payload if any
                 try:
-                    if (
-                        isinstance(req_state, OmniRequestState)
-                        and req_state.mm_accumulated is not None
-                    ):
+                    if isinstance(req_state, OmniRequestState) and req_state.mm_accumulated is not None:
                         mm_key = req_state.mm_type or "latents"
                         if not hasattr(ro, "multimodal_output"):
                             setattr(ro, "multimodal_output", {})
                         ro.multimodal_output[mm_key] = req_state.mm_accumulated
                 except Exception:
-                    logger.exception(
-                        "Error attaching multimodal payload in process_outputs"
-                    )
+                    logger.exception("Error attaching multimodal payload in process_outputs")
                 if req_state.queue is not None:
                     req_state.queue.put(ro)
                 else:
@@ -335,9 +304,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                     self.parent_requests.pop(parent_req.request_id, None)
                 if not eco.finished:
                     reqs_to_abort.append(req_id)
-                self._update_stats_from_finished(
-                    req_state, finish_reason, iteration_stats
-                )
+                self._update_stats_from_finished(req_state, finish_reason, iteration_stats)
                 if self.tracer:
                     self.do_tracing(eco, req_state, iteration_stats)
                 # Cleanup per-request mm state
@@ -352,9 +319,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
 
     # ---- routing helpers ----
     def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
-        output_type = (
-            getattr(eco, "output_type", self.engine_core_output_type) or ""
-        ).lower()
+        output_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
 
         # Custom handler first (if registered)
         if output_type in self.output_handlers:
@@ -385,9 +350,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
     def _process_image_output(self, eco: EngineCoreOutput) -> None:
         """Ensure image tensors are surfaced via pooling_output for vLLM."""
         if eco.pooling_output is None:
-            tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("image", "images", "pixel_values", "pixels")
-            )
+            tensor = self._extract_from_multimodal_outputs(eco, keys=("image", "images", "pixel_values", "pixels"))
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -414,9 +377,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
     def _process_latents_output(self, eco: EngineCoreOutput) -> None:
         """Ensure latent tensors are surfaced via pooling_output."""
         if eco.pooling_output is None:
-            tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("latent", "latents", "z", "posterior")
-            )
+            tensor = self._extract_from_multimodal_outputs(eco, keys=("latent", "latents", "z", "posterior"))
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -444,9 +405,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             except Exception:
                 pass
 
-    def _extract_from_multimodal_outputs(
-        self, eco: EngineCoreOutput, keys: tuple[str, ...]
-    ) -> Optional[torch.Tensor]:
+    def _extract_from_multimodal_outputs(self, eco: EngineCoreOutput, keys: tuple[str, ...]) -> Optional[torch.Tensor]:
         mm = getattr(eco, "multimodal_outputs", None)
         if not isinstance(mm, dict):
             return None

vllm_omni/engine/processor.py

@@ -3,10 +3,10 @@ from collections.abc import Mapping
 from typing import Any, Optional, Union
 
 import torch
-
 from vllm.config import VllmConfig
 from vllm.inputs import ProcessorInputs, PromptType
 from vllm.inputs.parse import split_enc_dec_inputs
+from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
 from vllm.multimodal.inputs import MultiModalFeatureSpec
@@ -17,6 +17,7 @@ from vllm.sampling_params import SamplingParams
 from vllm.transformers_utils.tokenizer import AnyTokenizer
 from vllm.utils import length_from_prompt_token_ids_or_embeds
 from vllm.v1.engine.processor import Processor
+
 from vllm_omni.engine import (
     AdditionalInformationEntry,
     AdditionalInformationPayload,
@@ -25,6 +26,8 @@ from vllm_omni.engine import (
 )
 from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 
+logger = init_logger(__name__)
+
 
 class OmniProcessor(Processor):
     @staticmethod
@@ -75,20 +78,14 @@ class OmniProcessor(Processor):
         priority: int = 0,
         data_parallel_rank: Optional[int] = None,
     ) -> tuple[Optional[str], OmniEngineCoreRequest]:
-
         # TODO(woosuk): Support pooling models.
         # TODO(woosuk): Support encoder-decoder models.
         self._validate_lora(lora_request)
         self._validate_params(params)
 
         data_parallel_size = self.vllm_config.parallel_config.data_parallel_size
-        if data_parallel_rank is not None and not (
-            0 <= data_parallel_rank < data_parallel_size
-        ):
-            raise ValueError(
-                f"data_parallel_rank {data_parallel_rank} "
-                f"is out of range [0, {data_parallel_size})."
-            )
+        if data_parallel_rank is not None and not (0 <= data_parallel_rank < data_parallel_size):
+            raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {data_parallel_size}).")
 
         if arrival_time is None:
             arrival_time = time.time()
@@ -139,19 +136,9 @@ class OmniProcessor(Processor):
         # discriminated unions of TypedDicts, because of how it handles
         # inheritance of TypedDict. If we explicitly extract the items we want
         # we can avoid type errors from using `dict.get` later in the method.
-        prompt_str: Optional[str] = (
-            None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt")
-        )
-        prompt_token_ids = (
-            decoder_inputs["prompt_token_ids"]
-            if decoder_inputs["type"] != "embeds"
-            else None
-        )
-        prompt_embeds = (
-            decoder_inputs["prompt_embeds"]
-            if decoder_inputs["type"] == "embeds"
-            else None
-        )
+        prompt_str: Optional[str] = None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt")
+        prompt_token_ids = decoder_inputs["prompt_token_ids"] if decoder_inputs["type"] != "embeds" else None
+        prompt_embeds = decoder_inputs["prompt_embeds"] if decoder_inputs["type"] == "embeds" else None
 
         sampling_params = None
         pooling_params = None
@@ -160,13 +147,9 @@ class OmniProcessor(Processor):
             sampling_params = params.clone()
             # If unset max tokens, then generate up to the max_model_len.
             if sampling_params.max_tokens is None:
-                seq_len = length_from_prompt_token_ids_or_embeds(
-                    prompt_token_ids, prompt_embeds
-                )
+                seq_len = length_from_prompt_token_ids_or_embeds(prompt_token_ids, prompt_embeds)
                 sampling_params.max_tokens = self.model_config.max_model_len - seq_len
-            sampling_params.update_from_generation_config(
-                self.generation_config_fields, eos_token_id
-            )
+            sampling_params.update_from_generation_config(self.generation_config_fields, eos_token_id)
             if self.tokenizer is not None:
                 sampling_params.update_from_tokenizer(self.tokenizer)
         else:
@@ -201,12 +184,9 @@ class OmniProcessor(Processor):
         prompt_embeds_payload: Optional[PromptEmbedsPayload] = None
         additional_information_payload: Optional[AdditionalInformationPayload] = None
         if "prompt_embeds" in decoder_inputs:  # type: ignore[operator]
-
             pe: torch.Tensor = decoder_inputs["prompt_embeds"]  # type: ignore[index]
             if pe.ndim != 2:
-                raise ValueError(
-                    "prompt_embeds must be of shape (seq_len, hidden_size)"
-                )
+                raise ValueError("prompt_embeds must be of shape (seq_len, hidden_size)")
             # Move to CPU and ensure contiguous memory for stable serialization
             pe_cpu = pe.detach().to("cpu").contiguous()
             seq_len, hidden_size = pe_cpu.shape
@@ -233,13 +213,9 @@ class OmniProcessor(Processor):
                 elif isinstance(value, list):
                     entry = AdditionalInformationEntry(list_data=value)
                 else:
-                    raise ValueError(
-                        "additional_information values must be Tensor or list"
-                    )
+                    raise ValueError("additional_information values must be Tensor or list")
                 entries[key] = entry
-            additional_information_payload = AdditionalInformationPayload(
-                entries=entries
-            )
+            additional_information_payload = AdditionalInformationPayload(entries=entries)
 
         return prompt_str, OmniEngineCoreRequest(
             request_id=request_id,

vllm_omni/entrypoints/async_omni_llm.py

@@ -0,0 +1,742 @@
+import asyncio
+import multiprocessing as mp
+import os
+import socket
+import sys
+import time
+from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Optional, Union
+
+import torch
+
+# External library imports (vLLM)
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.engine.protocol import EngineClient
+from vllm.inputs import PromptType
+from vllm.inputs.preprocess import InputPreprocessor
+from vllm.logger import init_logger
+from vllm.lora.request import LoRARequest
+from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
+from vllm.sampling_params import SamplingParams
+from vllm.tracing import init_tracer
+from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
+from vllm.transformers_utils.tokenizer import AnyTokenizer, init_tokenizer_from_configs
+from vllm.usage.usage_lib import UsageContext
+from vllm.utils import Device, deprecate_kwargs
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.engine.core_client import EngineCoreClient
+from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.executor.abstract import Executor
+from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager
+
+# Internal imports (our code)
+from vllm_omni.config import OmniModelConfig
+from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
+from vllm_omni.engine.output_processor import MultimodalOutputProcessor
+from vllm_omni.engine.processor import OmniProcessor
+from vllm_omni.entrypoints.log_utils import (
+    OrchestratorMetrics,
+    configure_orchestrator_logger,
+    init_stats_paths,
+    remove_old_logs,
+)
+from vllm_omni.entrypoints.omni_stage import OmniStage
+from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
+from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
+from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
+from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
+from vllm_omni.outputs import OmniRequestOutput
+
+logger = init_logger(__name__)
+
+
+class AsyncOmniLLM(EngineClient):
+    def __init__(
+        self,
+        model: str,
+        stage_configs_path: Optional[str] = None,
+        log_stats: bool = False,
+        log_file: Optional[str] = None,
+        init_sleep_seconds: int = 20,
+        shm_threshold_bytes: int = 65536,
+        batch_timeout: int = 10,
+        init_timeout: int = 300,
+        **kwargs,
+    ):
+        self.batch_timeout = batch_timeout
+        self._enable_stats: bool = bool(log_stats)
+
+        if stage_configs_path is None:
+            self.stage_configs = load_stage_configs_from_model(model)
+        else:
+            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
+
+        self.stage_list: list[OmniStage] = []
+        self.default_sampling_params_list: list[SamplingParams] = []
+        # Optional file handler for orchestrator
+        self._log_file = log_file
+        if self._log_file:
+            remove_old_logs(self._log_file, len(self.stage_list))
+            configure_orchestrator_logger(logger, self._log_file)
+
+        self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
+        self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
+
+    def _initialize_stages(
+        self,
+        model: str,
+        init_sleep_seconds: int,
+        shm_threshold_bytes: int,
+        init_timeout: int,
+    ) -> None:
+        self.stage_list: list[OmniStage] = []
+
+        # Build OmniStage instances in parallel, preserve original order
+        def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
+            idx, cfg = idx_cfg
+            return idx, OmniStage(cfg)
+
+        with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor:
+            futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)]
+            results: list[tuple[int, OmniStage]] = []
+            for fut in as_completed(futures):
+                results.append(fut.result())
+        results.sort(key=lambda x: x[0])
+        self.stage_list = [st for _, st in results]
+        self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
+        logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
+
+        self._ctx = mp.get_context("spawn")
+        self._stage_in_queues: list[mp.Queue] = []
+        self._stage_out_queues: list[mp.Queue] = []
+        self._init_sleep_seconds = max(0, int(init_sleep_seconds))
+        self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
+        self._start_stage_processes(model)
+        # Wait for all stages to report readiness before seeding
+        self._stages_ready: set[int] = set()
+        self._wait_for_stages_ready(timeout=init_timeout)
+
+    def _start_stage_processes(self, model: str) -> None:
+        for stage_id, stage in enumerate(self.stage_list):
+            # Use unbounded queues to avoid deadlock when seeding many requests
+            in_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            self._stage_in_queues.append(in_q)
+            self._stage_out_queues.append(out_q)
+
+            # Attach queues and start Stage-owned worker process
+            stage.attach_queues(in_q, out_q)
+            stage.init_stage_worker(
+                model,
+                is_async=True,
+                log_file=self._log_file,
+                shm_threshold_bytes=self._shm_threshold_bytes,
+                ctx=self._ctx,
+                batch_timeout=self.batch_timeout,
+            )
+            logger.debug("[Orchestrator] Stage-%s process started", stage_id)
+            time.sleep(self._init_sleep_seconds)
+
+    def close(self) -> None:
+        for q in self._stage_in_queues:
+            try:
+                q.put_nowait(None)
+            except Exception as e:
+                logger.warning(
+                    "[Orchestrator] Failed to send shutdown signal to \
+                        stage input queue: %s",
+                    e,
+                )
+        for stage in self.stage_list:
+            try:
+                stage.stop_stage_worker()
+            except Exception as e:
+                logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
+
+    def __del__(self) -> None:  # best-effort
+        try:
+            self.close()
+        except Exception as e:
+            logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
+
+    def shutdown(self):
+        """Shutdown, cleaning up the background proc and IPC."""
+        try:
+            self.close()
+        except Exception as e:
+            logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
+
+    async def generate(
+        self,
+        prompt: PromptType,
+        request_id: str,
+        sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
+        lora_request: Optional[LoRARequest] = None,
+        trace_headers: Optional[Mapping[str, str]] = None,
+        priority: int = 0,
+        data_parallel_rank: Optional[int] = None,
+    ) -> AsyncGenerator[OmniRequestOutput, None]:
+        logger.debug("[Orchestrator] generate() called")
+        if sampling_params_list is None:
+            sampling_params_list = self.default_sampling_params_list
+        if len(sampling_params_list) != len(self.stage_list):
+            raise ValueError(
+                f"Expected {len(self.stage_list)} sampling params, \
+                got {len(sampling_params_list)}"
+            )
+
+        # Orchestrator keeps stage objects for input derivation
+        num_stages = len(self.stage_list)
+        # Track per-request start time for end-to-end timing
+        _req_start_ts: dict[int, float] = {}
+        _wall_start_ts: float = time.time()
+        # _last_finish_ts: float = _wall_start_ts
+
+        # Determine the final stage for E2E stats (highest stage_id with
+        # final_output=True; fallback to last stage)
+        final_stage_id_for_e2e = -1
+        try:
+            for _sid, _st in enumerate(self.stage_list):
+                if getattr(_st, "final_output", False):
+                    final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid)
+            if final_stage_id_for_e2e < 0:
+                final_stage_id_for_e2e = len(self.stage_list) - 1
+        except Exception as e:
+            logger.debug(
+                "[Orchestrator] Failed to determine final stage for E2E; \
+                    falling back to last: %s",
+                e,
+                exc_info=True,
+            )
+            final_stage_id_for_e2e = len(self.stage_list) - 1
+        # Metrics/aggregation helper
+        metrics = OrchestratorMetrics(
+            num_stages,
+            self._enable_stats,
+            self._stats_file,
+            self._overall_stats_file,
+            _wall_start_ts,
+        )
+        # Seed stage-0 queue with all requests
+        logger.debug("[Orchestrator] Seeding request into stage-0")
+        # Mark first input time for stage-0
+        metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
+
+        sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
+        task = {
+            "request_id": request_id,
+            "engine_inputs": prompt,
+            "sampling_params": sp0,
+        }
+        self.stage_list[0].submit(task)
+        _req_start_ts[request_id] = time.time()
+        logger.debug("[Orchestrator] Enqueued request %s to stage-0", request_id)
+
+        logger.debug("[Orchestrator] Entering scheduling loop: stages=%d", num_stages)
+        finished = False
+        while not finished:
+            made_progress = False
+            for stage_id, stage in enumerate(self.stage_list):
+                result = stage.try_collect()
+                if result is None:
+                    continue
+
+                made_progress = True
+                req_id = result.get("request_id")
+                if "error" in result:
+                    logger.error(
+                        "Stage %s error on request %s: %s",
+                        stage_id,
+                        req_id,
+                        result["error"],
+                    )
+                    continue
+
+                if result.get("type") == "stage_ready":
+                    # Only happens when stage is initialized slower than expected,
+                    # so we wait for a short time and try again
+                    time.sleep(0.05)
+                    continue
+
+                engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
+                # Mark last output time for this stage whenever we receive outputs
+                metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
+                try:
+                    _m = result.get("metrics")
+                    if _m is not None:
+                        metrics.on_stage_metrics(stage_id, req_id, _m)
+                except Exception as e:
+                    logger.exception(
+                        "[Orchestrator] Failed to process metrics for stage %s, \
+                            req %s: %s",
+                        stage_id,
+                        req_id,
+                        e,
+                    )
+                logger.debug(
+                    "[Orchestrator] Stage-%s completed request %s; \
+                        forwarding or finalizing",
+                    stage_id,
+                    req_id,
+                )
+                stage.set_engine_outputs(engine_outputs)
+
+                if getattr(stage, "final_output", False):
+                    logger.debug(
+                        "[Orchestrator] Request %s finalized at stage-%s",
+                        req_id,
+                        stage_id,
+                    )
+
+                    # End-to-end timing and time-per-token for final output
+                    # (only once per request at the designated final stage)
+                    try:
+                        rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
+                        if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
+                            metrics.on_finalize_request(
+                                stage_id,
+                                req_id,
+                                engine_outputs,
+                                _req_start_ts.get(req_id, _wall_start_ts),
+                            )
+                    except Exception as e:
+                        logger.exception(
+                            "[Orchestrator] Finalize request handling error for \
+                                req %s at stage %s: %s",
+                            req_id,
+                            stage_id,
+                            e,
+                        )
+
+                    if isinstance(engine_outputs, list):
+                        engine_outputs = engine_outputs[0]
+                    yield OmniRequestOutput(
+                        stage_id=stage_id,
+                        final_output_type=stage.final_output_type,
+                        request_output=engine_outputs,
+                    )
+
+                next_stage_id = stage_id + 1
+                if next_stage_id < num_stages:
+                    next_stage: OmniStage = self.stage_list[next_stage_id]
+                    next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
+                    sp_next: SamplingParams = sampling_params_list[next_stage_id]
+                    try:
+                        # Measure transfer size and time (encode + enqueue)
+                        size_bytes = 0
+                        try:
+                            size_bytes = len(_set(next_inputs))
+                        except Exception:
+                            size_bytes = 0
+                        t0 = time.time()
+                        ipc_payload = _encode(
+                            next_inputs,
+                            getattr(self, "_shm_threshold_bytes", 65536),
+                            obj_key="engine_inputs",
+                            shm_key="engine_inputs_shm",
+                        )
+                        ipc_payload.update(
+                            {
+                                "request_id": req_id,
+                                "sampling_params": sp_next,
+                                "sent_ts": time.time(),
+                            }
+                        )
+                        self.stage_list[next_stage_id].submit(ipc_payload)
+                        t1 = time.time()
+                        tx_ms = (t1 - t0) * 1000.0
+                        metrics.on_forward(
+                            stage_id,
+                            next_stage_id,
+                            req_id,
+                            int(size_bytes),
+                            float(tx_ms),
+                            bool("engine_inputs_shm" in ipc_payload),
+                        )
+                    except Exception as e:
+                        logger.warning(
+                            "[Orchestrator] IPC encode failed for req %s: %s; \
+                                falling back to inline payload",
+                            req_id,
+                            e,
+                        )
+                        self.stage_list[next_stage_id].submit(
+                            {
+                                "request_id": req_id,
+                                "engine_inputs": next_inputs,
+                                "sampling_params": sp_next,
+                            }
+                        )
+                    logger.debug(
+                        "[Orchestrator] Forwarded request %s to stage-%s",
+                        req_id,
+                        next_stage_id,
+                    )
+                else:
+                    finished = True
+                    logger.debug("[Orchestrator] Request %s fully completed", req_id)
+
+            if not made_progress:
+                time.sleep(0.005)
+
+        logger.debug("[Orchestrator] All requests completed")
+
+        # Summarize and print stats
+        try:
+            summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
+            logger.info("[Summary] %s", summary)
+        except Exception as e:
+            logger.exception("[Orchestrator] Failed to build/log summary: %s", e)
+
+    def _wait_for_stages_ready(self, timeout: int = 120) -> None:
+        deadline = time.time() + max(0, int(timeout))
+        num_stages = len(self.stage_list)
+        while len(self._stages_ready) < num_stages and time.time() < deadline:
+            progressed = False
+            for stage_id, stage in enumerate(self.stage_list):
+                if stage_id in self._stages_ready:
+                    continue
+                result = stage.try_collect()
+                if result is None:
+                    continue
+                progressed = True
+                if result.get("type") == "stage_ready":
+                    self._stages_ready.add(stage_id)
+                    # Store vllm_config received from worker process
+                    vllm_config = result.get("vllm_config")
+                    if vllm_config is not None:
+                        stage.set_vllm_config(vllm_config)
+                    tokenizer = result.get("tokenizer")
+                    if tokenizer is not None:
+                        stage.set_tokenizer(tokenizer)
+                    # input_preprocessor = result.get("input_preprocessor")
+                    # if input_preprocessor is not None:
+                    #     stage.set_input_preprocessor(input_preprocessor)
+                    is_tracing_enabled = result.get("is_tracing_enabled")
+                    if is_tracing_enabled is not None:
+                        stage.set_is_tracing_enabled(is_tracing_enabled)
+                    logger.debug("[Orchestrator] Stage-%s reported ready", stage_id)
+                else:
+                    # No user data should arrive before seeding; ignore other messages
+                    pass
+            if not progressed:
+                time.sleep(0.01)
+        if len(self._stages_ready) < num_stages:
+            not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
+            logger.warning(
+                "[Orchestrator] Initialization timeout: only %s/%s stages are \
+                    ready; not ready: %s",
+                len(self._stages_ready),
+                num_stages,
+                not_ready,
+            )
+            # Provide actionable suggestions before shutdown
+            try:
+                suggestions = [
+                    "Verify GPU/device assignment in config (runtime.devices) is \
+                        correct.",
+                    "Check GPU/host memory availability; reduce model or batch size if needed.",  # noqa: E501
+                    "Check model weights path and network reachability (if loading remotely).",  # noqa: E501
+                    "Increase initialization wait time (init_sleep_seconds or \
+                        call-site timeout).",
+                ]
+                if getattr(self, "_log_file", None):
+                    suggestions.append(
+                        f"Inspect per-stage log files for details: \
+                            {self._log_file}.stage<id>.log"
+                    )
+                logger.error(
+                    "[Orchestrator] Stage initialization failed, shutting down. \
+                        Suggestions:\n- %s",
+                    "\n- ".join(suggestions),
+                )
+            except Exception:
+                # Best-effort logging of suggestions
+                logger.error(
+                    "[Orchestrator] Stage initialization failed and an error \
+                        occurred while logging suggestions",
+                )
+
+            # Attempt graceful shutdown of all stages before exiting
+            try:
+                self.close()
+            except Exception:
+                pass
+
+            # Terminate the current process with non-zero exit code
+            try:
+                sys.exit(1)
+            except SystemExit:
+                raise
+            except Exception:
+                os._exit(1)
+
+    @property
+    def is_running(self) -> bool:
+        # Is None before the loop is started.
+        return len(self._stage_in_queues) > 0
+
+    @property
+    def is_stopped(self) -> bool:
+        return self.errored
+
+    @property
+    def errored(self) -> bool:
+        return not self.is_running
+
+    @property
+    def dead_error(self) -> BaseException:
+        return EngineDeadError()
+
+    async def abort(self, request_id: Union[str, Iterable[str]]) -> None:
+        pass
+
+    async def get_vllm_config(self) -> VllmConfig:
+        for stage in self.stage_list:
+            if stage.is_comprehension:
+                # Use the vllm_config received from worker process
+                if stage.vllm_config is not None:
+                    return stage.vllm_config
+        return None
+
+    async def get_model_config(self) -> OmniModelConfig:
+        for stage in self.stage_list:
+            if stage.is_comprehension:
+                # Use the vllm_config received from worker process
+                if stage.vllm_config is not None:
+                    return stage.vllm_config.model_config
+        return None
+
+    async def get_input_preprocessor(self) -> InputPreprocessor:
+        # for stage in self.stage_list:
+        #     if stage.is_comprehension:
+        #         return stage.input_preprocessor
+        return None
+
+    async def get_tokenizer(self) -> AnyTokenizer:
+        for stage in self.stage_list:
+            if stage.is_comprehension:
+                return stage.tokenizer
+        return None
+
+    async def is_tracing_enabled(self) -> bool:
+        for stage in self.stage_list:
+            if stage.is_comprehension:
+                return stage.is_tracing_enabled
+        return False
+
+    async def do_log_stats(self) -> None:
+        pass
+
+    async def check_health(self) -> None:
+        pass
+
+    async def reset_mm_cache(self) -> None:
+        pass
+
+    async def reset_prefix_cache(self, device: Optional[Device] = None) -> None:
+        pass
+
+    async def sleep(self, level: int = 1) -> None:
+        pass
+
+    async def wake_up(self, tags: Optional[list[str]] = None) -> None:
+        pass
+
+    async def is_sleeping(self) -> bool:
+        """Check whether the engine is sleeping"""
+        return False
+
+    async def add_lora(self, lora_request: LoRARequest) -> bool:
+        """Load a new LoRA adapter into the engine for future requests."""
+        return False
+
+    async def encode(
+        self,
+        *args,
+        **kwargs,
+    ):
+        """Generate outputs for a request from a pooling model."""
+        raise NotImplementedError("encode() is not implemented for AsyncOmniLLM")
+
+    async def start_profile(self) -> None:
+        raise NotImplementedError("start_profile() is not implemented for AsyncOmniLLM")
+
+    async def stop_profile(self) -> None:
+        raise NotImplementedError("stop_profile() is not implemented for AsyncOmniLLM")
+
+
+class AsyncOmniStageLLM(AsyncLLM):
+    def __init__(
+        self,
+        engine_args: AsyncOmniEngineArgs,
+        vllm_config: VllmConfig,
+        executor_class: type[Executor],
+        log_stats: bool,
+        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
+        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
+        use_cached_outputs: bool = False,
+        log_requests: bool = True,
+        start_engine_loop: bool = True,
+        stat_loggers: Optional[list[StatLoggerFactory]] = None,
+        client_addresses: Optional[dict[str, str]] = None,
+        client_count: int = 1,
+        client_index: int = 0,
+    ) -> None:
+        """
+        Create an AsyncLLM.
+
+        Args:
+            vllm_config: global configuration.
+            executor_class: an Executor impl, e.g. MultiprocExecutor.
+            log_stats: Whether to log stats.
+            usage_context: Usage context of the LLM.
+            mm_registry: Multi-modal registry.
+            use_cached_outputs: Whether to use cached outputs.
+            log_requests: Whether to log requests.
+            start_engine_loop: Whether to start the engine loop.
+            stat_loggers: customized stat loggers for the engine.
+                If not provided, default stat loggers will be used.
+                PLEASE BE AWARE THAT STAT LOGGER IS NOT STABLE
+                IN V1, AND ITS BASE CLASS INTERFACE MIGHT CHANGE.
+
+        Returns:
+            None
+        """
+        if not envs.VLLM_USE_V1:
+            raise ValueError(
+                "Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. "
+                "This should not happen. As a workaround, try using "
+                "AsyncLLMEngine.from_vllm_config(...) or explicitly set "
+                "VLLM_USE_V1=0 or 1 and report this issue on Github."
+            )
+
+        # Ensure we can serialize custom transformer configs
+        maybe_register_config_serialize_by_value()
+
+        self.model_config = vllm_config.model_config
+        self.vllm_config = vllm_config
+        self.observability_config = vllm_config.observability_config
+        self.log_requests = log_requests
+
+        self.log_stats = log_stats or (stat_loggers is not None)
+        if not log_stats and stat_loggers is not None:
+            logger.info(
+                "AsyncLLM created with log_stats=False and non-empty custom logger list; "
+                "enabling logging without default stat loggers"
+            )
+
+        if self.model_config.skip_tokenizer_init:
+            self.tokenizer = None
+        else:
+            # Tokenizer (+ ensure liveness if running in another process).
+            self.tokenizer = init_tokenizer_from_configs(model_config=vllm_config.model_config)
+
+        # Processor (converts Inputs --> EngineCoreRequests).
+        self.processor = OmniProcessor(
+            vllm_config=vllm_config,
+            tokenizer=self.tokenizer,
+            mm_registry=mm_registry,
+        )
+
+        # OutputProcessor (converts EngineCoreOutputs --> RequestOutput).
+        self.output_processor = MultimodalOutputProcessor(
+            tokenizer=self.tokenizer,
+            log_stats=self.log_stats,
+            engine_core_output_type=engine_args.engine_output_type,
+        )
+        if self.observability_config.otlp_traces_endpoint is not None:
+            tracer = init_tracer("vllm.llm_engine", self.observability_config.otlp_traces_endpoint)
+            self.output_processor.tracer = tracer
+
+        # EngineCore (starts the engine in background process).
+        self.engine_core = EngineCoreClient.make_async_mp_client(
+            vllm_config=vllm_config,
+            executor_class=executor_class,
+            log_stats=self.log_stats,
+            client_addresses=client_addresses,
+            client_count=client_count,
+            client_index=client_index,
+        )
+
+        # Loggers.
+        self.logger_manager: Optional[StatLoggerManager] = None
+        if self.log_stats:
+            self.logger_manager = StatLoggerManager(
+                vllm_config=vllm_config,
+                engine_idxs=self.engine_core.engine_ranks_managed,
+                custom_stat_loggers=stat_loggers,
+                enable_default_loggers=log_stats,
+                client_count=client_count,
+            )
+            self.logger_manager.log_engine_initialized()
+
+        self.output_handler: Optional[asyncio.Task] = None
+        try:
+            # Start output handler eagerly if we are in the asyncio eventloop.
+            asyncio.get_running_loop()
+            self._run_output_handler()
+        except RuntimeError:
+            pass
+
+        if envs.VLLM_TORCH_PROFILER_DIR:
+            logger.info(
+                "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s",  # noqa: E501
+                envs.VLLM_TORCH_PROFILER_DIR,
+            )
+            worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm"
+            self.profiler = torch.profiler.profile(
+                activities=[
+                    torch.profiler.ProfilerActivity.CPU,
+                ],
+                with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
+                on_trace_ready=torch.profiler.tensorboard_trace_handler(
+                    envs.VLLM_TORCH_PROFILER_DIR, worker_name=worker_name, use_gzip=True
+                ),
+            )
+        else:
+            self.profiler = None
+
+    @classmethod
+    @deprecate_kwargs(
+        "disable_log_requests",
+        additional_message=("This argument will have no effect. Use `enable_log_requests` instead."),
+    )
+    def from_vllm_config(
+        cls,
+        vllm_config: VllmConfig,
+        engine_args: AsyncOmniEngineArgs,
+        start_engine_loop: bool = True,
+        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
+        stat_loggers: Optional[list[StatLoggerFactory]] = None,
+        enable_log_requests: bool = False,
+        disable_log_stats: bool = False,
+        client_addresses: Optional[dict[str, str]] = None,
+        client_count: int = 1,
+        client_index: int = 0,
+        disable_log_requests: bool = True,  # Deprecated, will be removed
+    ) -> "AsyncLLM":
+        if not envs.VLLM_USE_V1:
+            raise ValueError(
+                "Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. "
+                "This should not happen. As a workaround, try using "
+                "AsyncLLMEngine.from_vllm_config(...) or explicitly set "
+                "VLLM_USE_V1=0 or 1 and report this issue on Github."
+            )
+
+        # Create the LLMEngine.
+        return cls(
+            vllm_config=vllm_config,
+            executor_class=Executor.get_class(vllm_config),
+            start_engine_loop=start_engine_loop,
+            stat_loggers=stat_loggers,
+            log_requests=enable_log_requests,
+            log_stats=not disable_log_stats,
+            usage_context=usage_context,
+            client_addresses=client_addresses,
+            client_count=client_count,
+            client_index=client_index,
+            engine_args=engine_args,
+        )

vllm_omni/entrypoints/cli/main.py

@@ -2,24 +2,55 @@
 CLI entry point for vLLM-omni that intercepts vLLM commands.
 """
 
+import importlib.metadata
 import sys
 
-from vllm_omni.entrypoints.cli.serve import OmniServeCommand
-
 
 def main():
     """Main CLI entry point that intercepts vLLM commands."""
     # Check if --omni flag is present
-    if "--omni" in sys.argv:
-        # Remove --omni flag and "serve" command, process with vLLM-omni
-        omni_args = [arg for arg in sys.argv[1:] if arg not in ["--omni", "serve"]]
-        omni_serve = OmniServeCommand()
-        omni_serve.run(omni_args)
-    else:
-        # Forward to original vLLM CLI
+    if "--omni" not in sys.argv:
         from vllm.entrypoints.cli.main import main as vllm_main
 
         vllm_main()
+        return
+    else:
+        from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup
+        from vllm.utils import FlexibleArgumentParser
+
+        import vllm_omni.entrypoints.cli.serve
+
+        CMD_MODULES = [
+            vllm_omni.entrypoints.cli.serve,
+        ]
+
+        cli_env_setup()
+
+        parser = FlexibleArgumentParser(
+            description="vLLM OMNI CLI",
+            epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"),
+        )
+        parser.add_argument(
+            "-v",
+            "--version",
+            action="version",
+            version=importlib.metadata.version("vllm_omni"),
+        )
+        subparsers = parser.add_subparsers(required=False, dest="subparser")
+        cmds = {}
+        for cmd_module in CMD_MODULES:
+            new_cmds = cmd_module.cmd_init()
+            for cmd in new_cmds:
+                cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
+                cmds[cmd.name] = cmd
+        args = parser.parse_args()
+        if args.subparser in cmds:
+            cmds[args.subparser].validate(args)
+
+        if hasattr(args, "dispatch_function"):
+            args.dispatch_function(args)
+        else:
+            parser.print_help()
 
 
 if __name__ == "__main__":

vllm_omni/entrypoints/cli/serve.py

@@ -3,93 +3,65 @@ Omni serve command for vLLM-omni.
 """
 
 import argparse
-from typing import List
+
+import uvloop
+from vllm.entrypoints.cli.types import CLISubcommand
+from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
+from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
+from vllm.logger import init_logger
+from vllm.utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.openai.api_server import omni_run_server
+
+logger = init_logger(__name__)
+
+DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni LLM
+completions via HTTP. Defaults to Qwen/Qwen2.5-Omni-7B if no model is specified.
+
+Search by using: `--help=<ConfigGroup>` to explore options by section (e.g.,
+--help=ModelConfig, --help=Frontend)
+  Use `--help=all` to show all available flags at once.
+"""
 
 
-class OmniServeCommand:
-    """Command handler for vLLM-omni serve command."""
+class OmniServeCommand(CLISubcommand):
+    """The `serve` subcommand for the vLLM CLI."""
 
-    def __init__(self):
-        self.parser = self._create_parser()
+    name = "serve"
 
-    def _create_parser(self) -> argparse.ArgumentParser:
-        """Create argument parser for omni serve command."""
-        parser = argparse.ArgumentParser(
-            description="vLLM-omni: Multi-modality models inference and serving"
+    @staticmethod
+    def cmd(args: argparse.Namespace) -> None:
+        # If model is specified in CLI (as positional arg), it takes precedence
+        if hasattr(args, "model_tag") and args.model_tag is not None:
+            args.model = args.model_tag
+
+        uvloop.run(omni_run_server(args))
+
+    def validate(self, args: argparse.Namespace) -> None:
+        validate_parsed_serve_args(args)
+
+    def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
+        serve_parser = subparsers.add_parser(
+            self.name,
+            description=DESCRIPTION,
+            usage="vllm serve [model_tag] --omni [options]",
         )
 
-        # Model arguments - make it optional with default
-        parser.add_argument(
-            "model",
-            nargs="?",
-            default="Qwen/Qwen3-0.6B",
-            help="Path to the model or model name (default: Qwen/Qwen3-0.6B)",
-        )
-
-        # Server arguments
-        parser.add_argument(
-            "--port",
-            type=int,
-            default=8000,
-            help="Port to run the server on",
-        )
-
-        parser.add_argument(
-            "--host",
-            type=str,
-            default="0.0.0.0",
-            help="Host to run the server on",
-        )
-
-        # Stage configuration arguments
-        parser.add_argument("--ar-stage", type=str, help="AR stage model path")
-
-        parser.add_argument("--dit-stage", type=str, help="DiT stage model path")
-
-        parser.add_argument(
-            "--dit-steps",
-            type=int,
-            default=50,
-            help="Number of DiT inference steps",
-        )
-
-        parser.add_argument(
-            "--dit-guidance-scale",
-            type=float,
-            default=7.5,
-            help="DiT guidance scale",
-        )
-
-        parser.add_argument(
-            "--use-diffusers",
+        serve_parser = make_arg_parser(serve_parser)
+        serve_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name)
+        serve_parser.add_argument(
+            "--omni",
             action="store_true",
-            help="Use diffusers pipeline for DiT stage",
+            help="Enable vLLM-omni mode",
         )
-
-        # Other arguments
-        parser.add_argument(
-            "--log-stats", action="store_true", help="Enable logging statistics"
+        serve_parser.add_argument(
+            "--stage-configs-path",
+            type=str,
+            default=None,
+            help="Path to the stage configs file. If not specified, the stage configs will be loaded from the model.",
         )
+        return serve_parser
 
-        return parser
 
-    def run(self, args: List[str]) -> None:
-        """Run the omni serve command."""
-        pass
-
-    def _create_stage_configs(self, args) -> List:
-        """Create stage configurations based on arguments."""
-        pass
-
-    async def _start_server(self, omni_llm, args) -> None:
-        """Start the API server."""
-        try:
-            # Import here to avoid circular imports
-            from vllm_omni.entrypoints.api_server import run_server
-
-            await run_server(omni_llm_instance=omni_llm, host=args.host, port=args.port)
-        except KeyboardInterrupt:
-            print("\nShutting down server...")
-        except Exception as e:
-            print(f"Error starting server: {e}")
-            raise
+def cmd_init() -> list[CLISubcommand]:
+    return [OmniServeCommand()]

vllm_omni/entrypoints/log_utils.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional
-
 import logging
 import os
 import time
+from typing import Any
 
 from vllm_omni.entrypoints.stage_utils import append_jsonl as _append_jsonl
 
 
-def remove_old_logs(log_file: Optional[str], num_stages: int) -> None:
+def remove_old_logs(log_file: str | None, num_stages: int) -> None:
     try:
         if not log_file:
             return
@@ -49,7 +48,7 @@ def remove_old_logs(log_file: Optional[str], num_stages: int) -> None:
         pass
 
 
-def configure_orchestrator_logger(logger: logging.Logger, log_file: Optional[str]) -> None:
+def configure_orchestrator_logger(logger: logging.Logger, log_file: str | None) -> None:
     try:
         if not log_file:
             return
@@ -64,9 +63,9 @@ def configure_orchestrator_logger(logger: logging.Logger, log_file: Optional[str
         pass
 
 
-def init_stats_paths(enable_stats: bool, log_file: Optional[str]) -> tuple[Optional[str], Optional[str]]:
-    stats_file: Optional[str] = None
-    overall_file: Optional[str] = None
+def init_stats_paths(enable_stats: bool, log_file: str | None) -> tuple[str | None, str | None]:
+    stats_file: str | None = None
+    overall_file: str | None = None
     try:
         if enable_stats and log_file:
             stats_file = f"{log_file}.orchestrator.stats.jsonl"
@@ -77,7 +76,7 @@ def init_stats_paths(enable_stats: bool, log_file: Optional[str]) -> tuple[Optio
     return stats_file, overall_file
 
 
-def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None:
+def _safe_append_jsonl(path: str | None, record: dict[str, Any]) -> None:
     if not path:
         return
     try:
@@ -86,122 +85,204 @@ def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None:
         pass
 
 
-def log_transfer_tx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
-                    size_bytes: int, tx_time_ms: float, used_shm: bool) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "transfer_stats",
-        "from_stage": from_stage,
-        "to_stage": to_stage,
-        "request_id": request_id,
-        "size_bytes": int(size_bytes),
-        "tx_time_ms": float(tx_time_ms),
-        "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
-        "used_shm": bool(used_shm),
-    })
+def log_transfer_tx(
+    stats_file: str | None,
+    from_stage: int,
+    to_stage: int,
+    request_id: Any,
+    size_bytes: int,
+    tx_time_ms: float,
+    used_shm: bool,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "transfer_stats",
+            "from_stage": from_stage,
+            "to_stage": to_stage,
+            "request_id": request_id,
+            "size_bytes": int(size_bytes),
+            "tx_time_ms": float(tx_time_ms),
+            "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
+            "used_shm": bool(used_shm),
+        },
+    )
 
 
-def log_transfer_rx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
-                    rx_bytes: int, rx_decode_time_ms: float, in_flight_time_ms: float) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "transfer_rx_stats",
-        "from_stage": from_stage,
-        "to_stage": to_stage,
-        "request_id": request_id,
-        "rx_bytes": int(rx_bytes),
-        "rx_decode_time_ms": float(rx_decode_time_ms),
-        "in_flight_time_ms": float(in_flight_time_ms),
-        "rx_time_per_kb_ms": (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0,
-    })
+def log_transfer_rx(
+    stats_file: str | None,
+    from_stage: int,
+    to_stage: int,
+    request_id: Any,
+    rx_bytes: int,
+    rx_decode_time_ms: float,
+    in_flight_time_ms: float,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "transfer_rx_stats",
+            "from_stage": from_stage,
+            "to_stage": to_stage,
+            "request_id": request_id,
+            "rx_bytes": int(rx_bytes),
+            "rx_decode_time_ms": float(rx_decode_time_ms),
+            "in_flight_time_ms": float(in_flight_time_ms),
+            "rx_time_per_kb_ms": (
+                (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0
+            ),
+        },
+    )
 
 
-def log_transfer_total(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
-                       size_bytes: int, tx_time_ms: float, in_flight_time_ms: float,
-                       rx_decode_time_ms: float, total_time_ms: float) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "transfer_total_stats",
-        "from_stage": from_stage,
-        "to_stage": to_stage,
-        "request_id": request_id,
-        "size_bytes": int(size_bytes),
-        "tx_time_ms": float(tx_time_ms),
-        "in_flight_time_ms": float(in_flight_time_ms),
-        "rx_decode_time_ms": float(rx_decode_time_ms),
-        "total_time_ms": float(total_time_ms),
-        "total_time_per_kb_ms": float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0,
-    })
+def log_transfer_total(
+    stats_file: str | None,
+    from_stage: int,
+    to_stage: int,
+    request_id: Any,
+    size_bytes: int,
+    tx_time_ms: float,
+    in_flight_time_ms: float,
+    rx_decode_time_ms: float,
+    total_time_ms: float,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "transfer_total_stats",
+            "from_stage": from_stage,
+            "to_stage": to_stage,
+            "request_id": request_id,
+            "size_bytes": int(size_bytes),
+            "tx_time_ms": float(tx_time_ms),
+            "in_flight_time_ms": float(in_flight_time_ms),
+            "rx_decode_time_ms": float(rx_decode_time_ms),
+            "total_time_ms": float(total_time_ms),
+            "total_time_per_kb_ms": (
+                float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0
+            ),
+        },
+    )
 
 
-def log_orchestrator_e2e(stats_file: Optional[str], request_id: Any, final_stage_id: int,
-                         e2e_time_ms: float, num_tokens_out: int, e2e_time_per_token_ms: float) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "orchestrator_request_e2e",
-        "request_id": request_id,
-        "final_stage_id": final_stage_id,
-        "e2e_time_ms": float(e2e_time_ms),
-        "num_tokens_out": int(num_tokens_out),
-        "e2e_time_per_token_ms": float(e2e_time_per_token_ms),
-    })
+def log_orchestrator_e2e(
+    stats_file: str | None,
+    request_id: Any,
+    final_stage_id: int,
+    e2e_time_ms: float,
+    num_tokens_out: int,
+    e2e_time_per_token_ms: float,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "orchestrator_request_e2e",
+            "request_id": request_id,
+            "final_stage_id": final_stage_id,
+            "e2e_time_ms": float(e2e_time_ms),
+            "num_tokens_out": int(num_tokens_out),
+            "e2e_time_per_token_ms": float(e2e_time_per_token_ms),
+        },
+    )
 
 
-def log_orchestrator_summary(stats_file: Optional[str], summary: Dict[str, Any]) -> None:
+def log_orchestrator_summary(stats_file: str | None, summary: dict[str, Any]) -> None:
     _safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary})
 
 
-def log_overall_summary(overall_stats_file: Optional[str], summary: Dict[str, Any]) -> None:
+def log_overall_summary(overall_stats_file: str | None, summary: dict[str, Any]) -> None:
     _safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary})
 
 
-def log_overall_record(overall_stats_file: Optional[str], record: Dict[str, Any]) -> None:
+def log_overall_record(overall_stats_file: str | None, record: dict[str, Any]) -> None:
     _safe_append_jsonl(overall_stats_file, record)
 
 
-def log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any, batch_size: int,
-                            num_tokens_out: int, stage_gen_time_ms: float, tokens_per_s: float,
-                            rx_transfer_bytes: int, rx_decode_time_ms: float, rx_mbps: float) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "stage_request_stats",
-        "stage_id": stage_id,
-        "request_id": request_id,
-        "batch_size": int(batch_size),
-        "num_tokens_out": int(num_tokens_out),
-        "stage_gen_time_ms": float(stage_gen_time_ms),
-        "tokens_per_s": float(tokens_per_s),
-        "rx_transfer_bytes": int(rx_transfer_bytes),
-        "rx_decode_time_ms": float(rx_decode_time_ms),
-        "rx_mbps": float(rx_mbps),
-    })
+def log_stage_request_stats(
+    stats_file: str | None,
+    stage_id: int,
+    request_id: Any,
+    batch_size: int,
+    num_tokens_out: int,
+    stage_gen_time_ms: float,
+    tokens_per_s: float,
+    rx_transfer_bytes: int,
+    rx_decode_time_ms: float,
+    rx_mbps: float,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "stage_request_stats",
+            "stage_id": stage_id,
+            "request_id": request_id,
+            "batch_size": int(batch_size),
+            "num_tokens_out": int(num_tokens_out),
+            "stage_gen_time_ms": float(stage_gen_time_ms),
+            "tokens_per_s": float(tokens_per_s),
+            "rx_transfer_bytes": int(rx_transfer_bytes),
+            "rx_decode_time_ms": float(rx_decode_time_ms),
+            "rx_mbps": float(rx_mbps),
+        },
+    )
 
 
-def log_stage_running_avg(stats_file: Optional[str], stage_id: int, total_tokens: int,
-                          total_gen_time_ms: float, avg_tokens_per_s: float) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "stage_running_avg",
-        "stage_id": stage_id,
-        "total_tokens": int(total_tokens),
-        "total_gen_time_ms": float(total_gen_time_ms),
-        "avg_tokens_per_s": float(avg_tokens_per_s),
-    })
+def log_stage_running_avg(
+    stats_file: str | None,
+    stage_id: int,
+    total_tokens: int,
+    total_gen_time_ms: float,
+    avg_tokens_per_s: float,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "stage_running_avg",
+            "stage_id": stage_id,
+            "total_tokens": int(total_tokens),
+            "total_gen_time_ms": float(total_gen_time_ms),
+            "avg_tokens_per_s": float(avg_tokens_per_s),
+        },
+    )
 
 
-def log_stage_batch_stats(stats_file: Optional[str], stage_id: int, batch_size: int,
-                          batch_gen_time_ms: float, request_ids: List[Any]) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "stage_batch_stats",
-        "stage_id": stage_id,
-        "batch_size": int(batch_size),
-        "batch_gen_time_ms": float(batch_gen_time_ms),
-        "request_ids": list(request_ids),
-    })
+def log_stage_batch_stats(
+    stats_file: str | None,
+    stage_id: int,
+    batch_size: int,
+    batch_gen_time_ms: float,
+    request_ids: list[Any],
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "stage_batch_stats",
+            "stage_id": stage_id,
+            "batch_size": int(batch_size),
+            "batch_gen_time_ms": float(batch_gen_time_ms),
+            "request_ids": list(request_ids),
+        },
+    )
 
 
-def compute_and_log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any,
-                                       batch_size: int, engine_outputs: List[Any],
-                                       stage_gen_time_ms: float, rx_transfer_bytes: int,
-                                       rx_decode_time_ms: float) -> None:
+def compute_and_log_stage_request_stats(
+    stats_file: str | None,
+    stage_id: int,
+    request_id: Any,
+    batch_size: int,
+    engine_outputs: list[Any],
+    stage_gen_time_ms: float,
+    rx_transfer_bytes: int,
+    rx_decode_time_ms: float,
+) -> None:
     """Compute per-request metrics and log them in one call."""
     num_tokens = count_tokens_from_outputs(engine_outputs)
     tokens_per_s = (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
-    rx_mbps = (float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0) if rx_transfer_bytes > 0 else 0.0
+    rx_mbps = (
+        (float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0)
+        if rx_transfer_bytes > 0
+        else 0.0
+    )
     log_stage_request_stats(
         stats_file,
         stage_id,
@@ -218,13 +299,16 @@ def compute_and_log_stage_request_stats(stats_file: Optional[str], stage_id: int
 
 # ----------------- Aggregation helpers for orchestrator -----------------
 
-def record_stage_metrics(per_request: Dict[int, Dict[str, Any]],
-                         stage_req_counts: List[int],
-                         stage_total_time_ms: List[float],
-                         stage_total_tokens: List[int],
-                         stage_id: int,
-                         req_id: Any,
-                         metrics: Dict[str, Any]) -> None:
+
+def record_stage_metrics(
+    per_request: dict[int, dict[str, Any]],
+    stage_req_counts: list[int],
+    stage_total_time_ms: list[float],
+    stage_total_tokens: list[int],
+    stage_id: int,
+    req_id: Any,
+    metrics: dict[str, Any],
+) -> None:
     try:
         stage_req_counts[stage_id] += 1
         stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
@@ -239,23 +323,32 @@ def record_stage_metrics(per_request: Dict[int, Dict[str, Any]],
         pass
 
 
-def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
-                                 transfer_agg: Dict[tuple[int, int], Dict[str, float]],
-                                 per_request: Dict[int, Dict[str, Any]],
-                                 stage_id: int,
-                                 req_id: Any,
-                                 rx_bytes: float,
-                                 rx_ms: float,
-                                 in_flight_ms: float) -> Optional[tuple[int, float, float]]:
+def aggregate_rx_and_maybe_total(
+    transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
+    transfer_agg: dict[tuple[int, int], dict[str, float]],
+    per_request: dict[int, dict[str, Any]],
+    stage_id: int,
+    req_id: Any,
+    rx_bytes: float,
+    rx_ms: float,
+    in_flight_ms: float,
+) -> tuple[int, float, float] | None:
     try:
         # Update RX aggregates for (stage_id-1 -> stage_id)
         if stage_id > 0:
             key = (stage_id - 1, stage_id)
             agg = transfer_agg.get(key)
             if agg is None:
-                agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
-                       "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
-                       "sum_total_ms": 0.0, "total_count": 0.0}
+                agg = {
+                    "sum_bytes": 0.0,
+                    "sum_ms": 0.0,
+                    "count": 0.0,
+                    "sum_rx_bytes": 0.0,
+                    "sum_rx_ms": 0.0,
+                    "rx_count": 0.0,
+                    "sum_total_ms": 0.0,
+                    "total_count": 0.0,
+                }
                 transfer_agg[key] = agg
             agg["sum_rx_bytes"] += float(rx_bytes)
             agg["sum_rx_ms"] += float(rx_ms)
@@ -287,20 +380,29 @@ def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], D
         return None
 
 
-def record_sender_transfer_agg(transfer_agg: Dict[tuple[int, int], Dict[str, float]],
-                               transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
-                               from_stage: int,
-                               to_stage: int,
-                               req_id: Any,
-                               size_bytes: int,
-                               tx_ms: float) -> None:
+def record_sender_transfer_agg(
+    transfer_agg: dict[tuple[int, int], dict[str, float]],
+    transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
+    from_stage: int,
+    to_stage: int,
+    req_id: Any,
+    size_bytes: int,
+    tx_ms: float,
+) -> None:
     try:
         key = (from_stage, to_stage)
         agg = transfer_agg.get(key)
         if agg is None:
-            agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
-                   "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
-                   "sum_total_ms": 0.0, "total_count": 0.0}
+            agg = {
+                "sum_bytes": 0.0,
+                "sum_ms": 0.0,
+                "count": 0.0,
+                "sum_rx_bytes": 0.0,
+                "sum_rx_ms": 0.0,
+                "rx_count": 0.0,
+                "sum_total_ms": 0.0,
+                "total_count": 0.0,
+            }
             transfer_agg[key] = agg
         agg["sum_bytes"] += float(size_bytes)
         agg["sum_ms"] += float(tx_ms)
@@ -318,7 +420,7 @@ def record_sender_transfer_agg(transfer_agg: Dict[tuple[int, int], Dict[str, flo
         pass
 
 
-def count_tokens_from_outputs(engine_outputs: List[Any]) -> int:
+def count_tokens_from_outputs(engine_outputs: list[Any]) -> int:
     total = 0
     for _ro in engine_outputs:
         try:
@@ -332,29 +434,35 @@ def count_tokens_from_outputs(engine_outputs: List[Any]) -> int:
     return total
 
 
-def build_stage_summary(stage_req_counts: List[int],
-                        stage_total_tokens: List[int],
-                        stage_total_time_ms: List[float]) -> List[Dict[str, Any]]:
-    summary: List[Dict[str, Any]] = []
+def build_stage_summary(
+    stage_req_counts: list[int],
+    stage_total_tokens: list[int],
+    stage_total_time_ms: list[float],
+) -> list[dict[str, Any]]:
+    summary: list[dict[str, Any]] = []
     for sid in range(len(stage_req_counts)):
         reqs = stage_req_counts[sid]
         tokens = stage_total_tokens[sid]
         total_ms = float(stage_total_time_ms[sid])
         avg_req = (total_ms / reqs) if reqs > 0 else 0.0
         avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
-        summary.append({
-            "stage_id": sid,
-            "requests": int(reqs),
-            "tokens": int(tokens),
-            "total_time_ms": total_ms,
-            "avg_time_per_request_ms": avg_req,
-            "avg_tokens_per_s": avg_tok,
-        })
+        summary.append(
+            {
+                "stage_id": sid,
+                "requests": int(reqs),
+                "tokens": int(tokens),
+                "total_time_ms": total_ms,
+                "avg_time_per_request_ms": avg_req,
+                "avg_tokens_per_s": avg_tok,
+            }
+        )
     return summary
 
 
-def build_transfer_summary(transfer_agg: Dict[tuple[int, int], Dict[str, float]]) -> List[Dict[str, Any]]:
-    summary: List[Dict[str, Any]] = []
+def build_transfer_summary(
+    transfer_agg: dict[tuple[int, int], dict[str, float]],
+) -> list[dict[str, Any]]:
+    summary: list[dict[str, Any]] = []
     for (src, dst), agg in transfer_agg.items():
         sum_bytes = float(agg.get("sum_bytes", 0.0))
         sum_ms = float(agg.get("sum_ms", 0.0))
@@ -367,50 +475,57 @@ def build_transfer_summary(transfer_agg: Dict[tuple[int, int], Dict[str, float]]
         sum_total_ms = float(agg.get("sum_total_ms", 0.0))
         samples_total = int(agg.get("total_count", 0.0))
         total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
-        summary.append({
-            "from_stage": src,
-            "to_stage": dst,
-            "samples": samples,
-            "total_bytes": int(sum_bytes),
-            "total_time_ms": sum_ms,
-            "tx_mbps": tx_mbps,
-            "rx_samples": samples_rx,
-            "rx_total_bytes": int(sum_rx_bytes),
-            "rx_total_time_ms": sum_rx_ms,
-            "rx_mbps": rx_mbps,
-            "total_samples": samples_total,
-            "total_transfer_time_ms": sum_total_ms,
-            "total_mbps": total_mbps,
-        })
+        summary.append(
+            {
+                "from_stage": src,
+                "to_stage": dst,
+                "samples": samples,
+                "total_bytes": int(sum_bytes),
+                "total_time_ms": sum_ms,
+                "tx_mbps": tx_mbps,
+                "rx_samples": samples_rx,
+                "rx_total_bytes": int(sum_rx_bytes),
+                "rx_total_time_ms": sum_rx_ms,
+                "rx_mbps": rx_mbps,
+                "total_samples": samples_total,
+                "total_transfer_time_ms": sum_total_ms,
+                "total_mbps": total_mbps,
+            }
+        )
     return summary
 
 
 class OrchestratorMetrics:
-    def __init__(self, num_stages: int, enable_stats: bool,
-                 stats_file: Optional[str], overall_stats_file: Optional[str],
-                 wall_start_ts: float) -> None:
+    def __init__(
+        self,
+        num_stages: int,
+        enable_stats: bool,
+        stats_file: str | None,
+        overall_stats_file: str | None,
+        wall_start_ts: float,
+    ) -> None:
         self.num_stages = int(num_stages)
         self.enable_stats = bool(enable_stats)
         self.stats_file = stats_file
         self.overall_stats_file = overall_stats_file
-        self.stage_total_time_ms: List[float] = [0.0 for _ in range(self.num_stages)]
-        self.stage_total_tokens: List[int] = [0 for _ in range(self.num_stages)]
-        self.stage_req_counts: List[int] = [0 for _ in range(self.num_stages)]
-        self.transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {}
-        self.transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {}
+        self.stage_total_time_ms: list[float] = [0.0 for _ in range(self.num_stages)]
+        self.stage_total_tokens: list[int] = [0 for _ in range(self.num_stages)]
+        self.stage_req_counts: list[int] = [0 for _ in range(self.num_stages)]
+        self.transfer_agg: dict[tuple[int, int], dict[str, float]] = {}
+        self.transfer_edge_req: dict[tuple[int, int, int], dict[str, float]] = {}
         self.e2e_total_ms: float = 0.0
         self.e2e_total_tokens: int = 0
         self.e2e_count: int = 0
         self.e2e_done: set[int] = set()
-        self.per_request: Dict[int, Dict[str, Any]] = {}
+        self.per_request: dict[int, dict[str, Any]] = {}
         self.sum_per_request_transfer_ms: float = 0.0
         self.wall_start_ts: float = float(wall_start_ts)
         self.last_finish_ts: float = float(wall_start_ts)
-        self.stage_seen_batches: Dict[int, set] = {sid: set() for sid in range(self.num_stages)}
-        self.stage_first_ts: List[Optional[float]] = [None for _ in range(self.num_stages)]
-        self.stage_last_ts: List[Optional[float]] = [None for _ in range(self.num_stages)]
+        self.stage_seen_batches: dict[int, set] = {sid: set() for sid in range(self.num_stages)}
+        self.stage_first_ts: list[float | None] = [None for _ in range(self.num_stages)]
+        self.stage_last_ts: list[float | None] = [None for _ in range(self.num_stages)]
 
-    def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: Dict[str, Any]) -> None:
+    def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: dict[str, Any]) -> None:
         record_stage_metrics(
             self.per_request,
             self.stage_req_counts,
@@ -466,8 +581,15 @@ class OrchestratorMetrics:
                     float(total_ms_c),
                 )
 
-    def on_forward(self, from_stage: int, to_stage: int, req_id: Any,
-                   size_bytes: int, tx_ms: float, used_shm: bool) -> None:
+    def on_forward(
+        self,
+        from_stage: int,
+        to_stage: int,
+        req_id: Any,
+        size_bytes: int,
+        tx_ms: float,
+        used_shm: bool,
+    ) -> None:
         # Mark first input time for the destination stage if not set
         if self.stage_first_ts[to_stage] is None:
             self.stage_first_ts[to_stage] = time.time()
@@ -491,8 +613,7 @@ class OrchestratorMetrics:
             float(tx_ms),
         )
 
-    def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: List[Any],
-                            req_start_ts: float) -> None:
+    def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: list[Any], req_start_ts: float) -> None:
         try:
             rid_int = int(req_id)
         except Exception:
@@ -526,30 +647,34 @@ class OrchestratorMetrics:
             e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0
             log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt)
 
-    def build_and_log_summary(self, final_stage_id_for_e2e: int) -> Dict[str, Any]:
+    def build_and_log_summary(self, final_stage_id_for_e2e: int) -> dict[str, Any]:
         # Compute stage summary using wall time between first input and last output per stage
-        stage_summary: List[Dict[str, Any]] = []
+        stage_summary: list[dict[str, Any]] = []
         for sid in range(self.num_stages):
             first_ts = self.stage_first_ts[sid]
             last_ts = self.stage_last_ts[sid]
-            total_ms = (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0
+            total_ms = (
+                (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0
+            )
             reqs = self.stage_req_counts[sid]
             tokens = self.stage_total_tokens[sid]
             avg_req = (total_ms / reqs) if reqs > 0 else 0.0
             avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
-            stage_summary.append({
-                "stage_id": sid,
-                "requests": int(reqs),
-                "tokens": int(tokens),
-                "total_time_ms": float(total_ms),
-                "avg_time_per_request_ms": float(avg_req),
-                "avg_tokens_per_s": float(avg_tok),
-            })
+            stage_summary.append(
+                {
+                    "stage_id": sid,
+                    "requests": int(reqs),
+                    "tokens": int(tokens),
+                    "total_time_ms": float(total_ms),
+                    "avg_time_per_request_ms": float(avg_req),
+                    "avg_tokens_per_s": float(avg_tok),
+                }
+            )
         transfer_summary = build_transfer_summary(self.transfer_agg)
         e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
         e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0
         wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
-        summary: Dict[str, Any] = {
+        summary: dict[str, Any] = {
             "e2e_requests": int(self.e2e_count),
             "e2e_total_time_ms": float(wall_time_ms),
             "e2e_sum_time_ms": float(self.e2e_total_ms),
@@ -566,5 +691,3 @@ class OrchestratorMetrics:
         if self.enable_stats and self.overall_stats_file:
             log_overall_summary(self.overall_stats_file, summary)
         return summary
-
-

vllm_omni/entrypoints/omni_llm.py

@@ -1,11 +1,10 @@
-from typing import Any, Optional, Sequence, Union
-import logging
 import multiprocessing as mp
-from concurrent.futures import ThreadPoolExecutor, as_completed
 import os
-import queue
 import sys
 import time
+from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Optional, Union
 
 import cloudpickle
 from pydantic import ValidationError
@@ -26,19 +25,17 @@ from vllm.v1.engine.llm_engine import LLMEngine
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
-from vllm_omni.entrypoints.omni_stage import OmniStage
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model
 from vllm_omni.entrypoints.log_utils import (
-    remove_old_logs,
+    OrchestratorMetrics,
     configure_orchestrator_logger,
     init_stats_paths,
-    OrchestratorMetrics,
-)
-from vllm_omni.entrypoints.stage_utils import (
-    maybe_load_from_ipc as _load,
-    encode_for_ipc as _encode,
-    serialize_obj as _ser,
+    remove_old_logs,
 )
+from vllm_omni.entrypoints.omni_stage import OmniStage
+from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
+from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
+from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
+from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -52,23 +49,25 @@ class OmniLLM:
     - max_inflight=1 per stage (serial within a stage), but pipeline across stages
     """
 
-    def __init__(self, model: str,
-                 stage_configs=None,
-                 log_stats: bool = False,
-                 log_file: Optional[str] = None,
-                 init_sleep_seconds: int = 20,
-                 shm_threshold_bytes: int = 65536,
-                 batch_timeout: int = 10,
-                 init_timeout: int = 300,
-                 **kwargs):
+    def __init__(
+        self,
+        model: str,
+        stage_configs_path: Optional[str] = None,
+        log_stats: bool = False,
+        log_file: Optional[str] = None,
+        init_sleep_seconds: int = 20,
+        shm_threshold_bytes: int = 65536,
+        batch_timeout: int = 10,
+        init_timeout: int = 300,
+        **kwargs,
+    ):
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
         # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
-        if stage_configs is None:
+        if stage_configs_path is None:
             self.stage_configs = load_stage_configs_from_model(model)
         else:
-            self.stage_configs = stage_configs
-
+            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
 
         # Optional file handler for orchestrator
         self._log_file = log_file
@@ -79,8 +78,15 @@ class OmniLLM:
         self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
         self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
 
-    def _initialize_stages(self, model: str, init_sleep_seconds: int, shm_threshold_bytes: int, init_timeout: int) -> None:
+    def _initialize_stages(
+        self,
+        model: str,
+        init_sleep_seconds: int,
+        shm_threshold_bytes: int,
+        init_timeout: int,
+    ) -> None:
         self.stage_list: list[OmniStage] = []
+
         # Build OmniStage instances in parallel, preserve original order
         def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
             idx, cfg = idx_cfg
@@ -105,7 +111,6 @@ class OmniLLM:
         self._stages_ready: set[int] = set()
         self._wait_for_stages_ready(timeout=init_timeout)
 
-
     def _start_stage_processes(self, model: str) -> None:
         for stage_id, stage in enumerate(self.stage_list):
             # Use unbounded queues to avoid deadlock when seeding many requests
@@ -131,7 +136,10 @@ class OmniLLM:
             try:
                 q.put_nowait(None)
             except Exception as e:
-                logger.warning("[Orchestrator] Failed to send shutdown signal to stage input queue: %s", e)
+                logger.warning(
+                    "[Orchestrator] Failed to send shutdown signal to stage input queue: %s",
+                    e,
+                )
         for stage in self.stage_list:
             try:
                 stage.stop_stage_worker()
@@ -147,17 +155,26 @@ class OmniLLM:
     def generate(
         self,
         prompts: Union[PromptType, Sequence[PromptType]],
-        sampling_params_list: Optional[
-            Union[SamplingParams, Sequence[SamplingParams]]
-        ] = None,
+        sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
+    ) -> list[OmniRequestOutput]:
+        try:
+            return self._run_generation(prompts, sampling_params_list)
+        except Exception as e:
+            logger.exception("[Orchestrator] Failed to run generation: %s", e)
+            raise e
+        finally:
+            self.close()
+
+    def _run_generation(
+        self,
+        prompts: Union[PromptType, Sequence[PromptType]],
+        sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
     ) -> list[OmniRequestOutput]:
         logger.debug("[Orchestrator] generate() called")
         if sampling_params_list is None:
             raise ValueError("sampling_params_list is required for pipelined generation")
         if len(sampling_params_list) != len(self.stage_list):
-            raise ValueError(
-                f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}"
-            )
+            raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
 
         # Normalize prompts to a list for per-request iteration
         if not isinstance(prompts, (list, tuple)):
@@ -176,7 +193,6 @@ class OmniLLM:
         # Track per-request start time for end-to-end timing
         _req_start_ts: dict[int, float] = {}
         _wall_start_ts: float = time.time()
-        _last_finish_ts: float = _wall_start_ts
 
         # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
         final_stage_id_for_e2e = -1
@@ -187,10 +203,20 @@ class OmniLLM:
             if final_stage_id_for_e2e < 0:
                 final_stage_id_for_e2e = len(self.stage_list) - 1
         except Exception as e:
-            logger.debug("[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s", e, exc_info=True)
+            logger.debug(
+                "[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s",
+                e,
+                exc_info=True,
+            )
             final_stage_id_for_e2e = len(self.stage_list) - 1
         # Metrics/aggregation helper
-        metrics = OrchestratorMetrics(num_stages, self._enable_stats, self._stats_file, self._overall_stats_file, _wall_start_ts)
+        metrics = OrchestratorMetrics(
+            num_stages,
+            self._enable_stats,
+            self._stats_file,
+            self._overall_stats_file,
+            _wall_start_ts,
+        )
 
         # Seed stage-0 queue with all requests
         logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts))
@@ -214,7 +240,11 @@ class OmniLLM:
         completed_requests = 0
         total_requests = len(request_prompts)
 
-        logger.debug("[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", total_requests, num_stages)
+        logger.debug(
+            "[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d",
+            total_requests,
+            num_stages,
+        )
         while completed_requests < total_requests:
             made_progress = False
             for stage_id, stage in enumerate(self.stage_list):
@@ -225,11 +255,17 @@ class OmniLLM:
                 made_progress = True
                 req_id = result.get("request_id")
                 if "error" in result:
-                    logger.error("Stage %s error on request %s: %s", stage_id, req_id, result["error"])
+                    logger.error(
+                        "Stage %s error on request %s: %s",
+                        stage_id,
+                        req_id,
+                        result["error"],
+                    )
                     continue
 
                 if result.get("type") == "stage_ready":
-                    #Only happens when stage is initialized slower than expected, so we wait for a short time and try again
+                    # Only happens when stage is initialized slower than expected,
+                    # so we wait for a short time and try again
                     time.sleep(0.05)
                     continue
 
@@ -241,8 +277,17 @@ class OmniLLM:
                     if _m is not None:
                         metrics.on_stage_metrics(stage_id, req_id, _m)
                 except Exception as e:
-                    logger.exception("[Orchestrator] Failed to process metrics for stage %s, req %s: %s", stage_id, req_id, e)
-                logger.debug("[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", stage_id, req_id)
+                    logger.exception(
+                        "[Orchestrator] Failed to process metrics for stage %s, req %s: %s",
+                        stage_id,
+                        req_id,
+                        e,
+                    )
+                logger.debug(
+                    "[Orchestrator] Stage-%s completed request %s; forwarding or finalizing",
+                    stage_id,
+                    req_id,
+                )
                 stage.set_engine_outputs(engine_outputs)
 
                 if getattr(stage, "final_output", False):
@@ -253,15 +298,30 @@ class OmniLLM:
                             request_output=engine_outputs,
                         )
                     )
-                    logger.debug("[Orchestrator] Request %s finalized at stage-%s", req_id, stage_id)
+                    logger.debug(
+                        "[Orchestrator] Request %s finalized at stage-%s",
+                        req_id,
+                        stage_id,
+                    )
 
-                    # End-to-end timing and time-per-token for final output (only once per request at the designated final stage)
+                    # End-to-end timing and time-per-token for final output
+                    # (only once per request at the designated final stage)
                     try:
                         rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
                         if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
-                            metrics.on_finalize_request(stage_id, req_id, engine_outputs, _req_start_ts.get(req_id, _wall_start_ts))
+                            metrics.on_finalize_request(
+                                stage_id,
+                                req_id,
+                                engine_outputs,
+                                _req_start_ts.get(req_id, _wall_start_ts),
+                            )
                     except Exception as e:
-                        logger.exception("[Orchestrator] Finalize request handling error for req %s at stage %s: %s", req_id, stage_id, e)
+                        logger.exception(
+                            "[Orchestrator] Finalize request handling error for req %s at stage %s: %s",
+                            req_id,
+                            stage_id,
+                            e,
+                        )
 
                 next_stage_id = stage_id + 1
                 if next_stage_id < num_stages:
@@ -272,7 +332,7 @@ class OmniLLM:
                         # Measure transfer size and time (encode + enqueue)
                         size_bytes = 0
                         try:
-                            size_bytes = len(_ser(next_inputs))
+                            size_bytes = len(_set(next_inputs))
                         except Exception:
                             size_bytes = 0
                         t0 = time.time()
@@ -282,27 +342,51 @@ class OmniLLM:
                             obj_key="engine_inputs",
                             shm_key="engine_inputs_shm",
                         )
-                        ipc_payload.update({
-                            "request_id": req_id,
-                            "sampling_params": sp_next,
-                            "sent_ts": time.time(),
-                        })
+                        ipc_payload.update(
+                            {
+                                "request_id": req_id,
+                                "sampling_params": sp_next,
+                                "sent_ts": time.time(),
+                            }
+                        )
                         self.stage_list[next_stage_id].submit(ipc_payload)
                         t1 = time.time()
                         tx_ms = (t1 - t0) * 1000.0
-                        metrics.on_forward(stage_id, next_stage_id, req_id, int(size_bytes), float(tx_ms), bool("engine_inputs_shm" in ipc_payload))
+                        metrics.on_forward(
+                            stage_id,
+                            next_stage_id,
+                            req_id,
+                            int(size_bytes),
+                            float(tx_ms),
+                            bool("engine_inputs_shm" in ipc_payload),
+                        )
                     except Exception as e:
-                        logger.warning("[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload", req_id, e)
-                        self.stage_list[next_stage_id].submit({
-                            "request_id": req_id,
-                            "engine_inputs": next_inputs,
-                            "sampling_params": sp_next,
-                        })
-                    logger.debug("[Orchestrator] Forwarded request %s to stage-%s", req_id, next_stage_id)
+                        logger.warning(
+                            "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
+                            req_id,
+                            e,
+                        )
+                        self.stage_list[next_stage_id].submit(
+                            {
+                                "request_id": req_id,
+                                "engine_inputs": next_inputs,
+                                "sampling_params": sp_next,
+                            }
+                        )
+                    logger.debug(
+                        "[Orchestrator] Forwarded request %s to stage-%s",
+                        req_id,
+                        next_stage_id,
+                    )
                     remaining_by_stage[next_stage_id] += 1
                 else:
                     completed_requests += 1
-                    logger.debug("[Orchestrator] Request %s fully completed (%d/%d)", req_id, completed_requests, total_requests)
+                    logger.debug(
+                        "[Orchestrator] Request %s fully completed (%d/%d)",
+                        req_id,
+                        completed_requests,
+                        total_requests,
+                    )
 
             if not made_progress:
                 time.sleep(0.005)
@@ -341,7 +425,9 @@ class OmniLLM:
             not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
             logger.warning(
                 "[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s",
-                len(self._stages_ready), num_stages, not_ready,
+                len(self._stages_ready),
+                num_stages,
+                not_ready,
             )
             # Provide actionable suggestions before shutdown
             try:
@@ -352,9 +438,7 @@ class OmniLLM:
                     "Increase initialization wait time (init_sleep_seconds or call-site timeout).",
                 ]
                 if getattr(self, "_log_file", None):
-                    suggestions.append(
-                        f"Inspect per-stage log files for details: {self._log_file}.stage<id>.log"
-                    )
+                    suggestions.append(f"Inspect per-stage log files for details: {self._log_file}.stage<id>.log")
                 logger.error(
                     "[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s",
                     "\n- ".join(suggestions),
@@ -384,13 +468,9 @@ class OmniStageLLM(LLM):
     def __init__(
         self,
         model: str,
-        compilation_config: Optional[
-            Union[int, dict[str, Any], CompilationConfig]
-        ] = None,
+        compilation_config: Optional[Union[int, dict[str, Any], CompilationConfig]] = None,
         hf_overrides: Optional[HfOverrides] = None,
-        structured_outputs_config: Optional[
-            Union[dict[str, Any], StructuredOutputsConfig]
-        ] = None,
+        structured_outputs_config: Optional[Union[dict[str, Any], StructuredOutputsConfig]] = None,
         **kwargs,
     ):
         """LLM constructor."""
@@ -404,9 +484,7 @@ class OmniStageLLM(LLM):
             if isinstance(worker_cls, type):
                 kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)
 
-        if "kv_transfer_config" in kwargs and isinstance(
-            kwargs["kv_transfer_config"], dict
-        ):
+        if "kv_transfer_config" in kwargs and isinstance(kwargs["kv_transfer_config"], dict):
             from vllm.config.kv_transfer import KVTransferConfig
 
             raw_config_dict = kwargs["kv_transfer_config"]
@@ -414,8 +492,7 @@ class OmniStageLLM(LLM):
                 kwargs["kv_transfer_config"] = KVTransferConfig(**raw_config_dict)
             except ValidationError as e:
                 logger.error(
-                    "Failed to convert 'kv_transfer_config' dict to "
-                    "KVTransferConfig object. Dict: %s. Error: %s",
+                    "Failed to convert 'kv_transfer_config' dict to KVTransferConfig object. Dict: %s. Error: %s",
                     raw_config_dict,
                     e,
                 )
@@ -428,16 +505,10 @@ class OmniStageLLM(LLM):
 
         if compilation_config is not None:
             if isinstance(compilation_config, int):
-                compilation_config_instance = CompilationConfig(
-                    level=compilation_config
-                )
+                compilation_config_instance = CompilationConfig(level=compilation_config)
             elif isinstance(compilation_config, dict):
                 compilation_config_instance = CompilationConfig(
-                    **{
-                        k: v
-                        for k, v in compilation_config.items()
-                        if is_init_field(CompilationConfig, k)
-                    }
+                    **{k: v for k, v in compilation_config.items() if is_init_field(CompilationConfig, k)}
                 )
             else:
                 compilation_config_instance = compilation_config
@@ -447,11 +518,7 @@ class OmniStageLLM(LLM):
         if structured_outputs_config is not None:
             if isinstance(structured_outputs_config, dict):
                 structured_outputs_instance = StructuredOutputsConfig(
-                    **{
-                        k: v
-                        for k, v in structured_outputs_config.items()
-                        if is_init_field(StructuredOutputsConfig, k)
-                    }
+                    **{k: v for k, v in structured_outputs_config.items() if is_init_field(StructuredOutputsConfig, k)}
                 )
             else:
                 structured_outputs_instance = structured_outputs_config
@@ -467,9 +534,7 @@ class OmniStageLLM(LLM):
         )
 
         # Create the Engine (autoselects V0 vs V1)
-        self.llm_engine = LLMEngine.from_engine_args(
-            engine_args=engine_args, usage_context=UsageContext.LLM_CLASS
-        )
+        self.llm_engine = LLMEngine.from_engine_args(engine_args=engine_args, usage_context=UsageContext.LLM_CLASS)
         self.llm_engine.output_processor = MultimodalOutputProcessor(
             tokenizer=self.llm_engine.tokenizer,
             log_stats=self.llm_engine.log_stats,
@@ -491,6 +556,4 @@ class OmniStageLLM(LLM):
 
         # Load the Input/Output processor plugin if any
         io_processor_plugin = self.llm_engine.model_config.io_processor_plugin
-        self.io_processor = get_io_processor(
-            self.llm_engine.vllm_config, io_processor_plugin
-        )
+        self.io_processor = get_io_processor(self.llm_engine.vllm_config, io_processor_plugin)

vllm_omni/entrypoints/omni_stage.py

@@ -1,64 +1,71 @@
 """
 Stage manager for orchestrating multiple engines in vLLM-omni.
 
+Enhanced to encapsulate per-stage process lifecycle and worker logic
+(device setup, LLM init, batching, shared-memory IPC), while preserving
+the original input processing utilities for cross-stage data wiring.
+
 Enhanced to encapsulate per-stage process lifecycle and worker logic
 (device setup, LLM init, batching, shared-memory IPC), while preserving
 the original input processing utilities for cross-stage data wiring.
 """
 
+import asyncio
 import importlib
 import logging
 import multiprocessing as mp
-from typing import Any, Dict, List, Optional, Union
-
-from vllm_omni.entrypoints.stage_utils import (
-    _to_dict,
-    set_stage_gpu_devices,
-    maybe_load_from_ipc_with_metrics,
-    maybe_dump_to_shm,
-)
+from typing import Any, Optional, Union
 
 from vllm.inputs import TextPrompt
+from vllm.inputs.preprocess import InputPreprocessor
+from vllm.logger import init_logger
+from vllm.sampling_params import SamplingParams
+from vllm.transformers_utils.tokenizer import AnyTokenizer
+from vllm.usage.usage_lib import UsageContext
 from vllm.v1.engine import EngineCoreOutput
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.engine.llm_engine import LLMEngine
+
+from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
+from vllm_omni.entrypoints.stage_utils import (
+    _to_dict,
+    maybe_dump_to_shm,
+    maybe_load_from_ipc_with_metrics,
+    set_stage_gpu_devices,
+)
 from vllm_omni.inputs.data import OmniTokensPrompt
 
+logger = init_logger(__name__)
+
 
 class OmniStage:
     def __init__(self, stage_config):
         self.stage_config = stage_config
         self.engine = None
         self.async_engine = None
+        self.vllm_config = None
+        self.tokenizer = None
+        self.input_preprocessor = None
+        self.is_tracing_enabled = False
         self.stage_id = stage_config.stage_id
         self.engine_args = stage_config.engine_args
         self.model_stage = stage_config.engine_args.model_stage
-        if hasattr(stage_config, "engine_input_source"):
-            self.engine_input_source = stage_config.engine_input_source
-        else:
-            self.engine_input_source = []
+        self.engine_input_source = getattr(stage_config, "engine_input_source", [])
         self.engine_output_type = stage_config.engine_args.engine_output_type
         self.engine_outputs = None
+        self.is_comprehension = getattr(stage_config, "is_comprehension", False)
         if hasattr(stage_config, "custom_process_input_func"):
             # Import the module specified in the config (already a full module path)
-            module_path, func_name = stage_config.custom_process_input_func.rsplit(
-                ".", 1
-            )
+            module_path, func_name = stage_config.custom_process_input_func.rsplit(".", 1)
             module = importlib.import_module(module_path)
             self.custom_process_input_func = getattr(module, func_name)
         else:
             self.custom_process_input_func = None
 
-        if hasattr(stage_config, "final_output"):
-            self.final_output = stage_config.final_output
-        else:
-            self.final_output = False
-
-        if hasattr(stage_config, "final_output_type"):
-            self.final_output_type = stage_config.final_output_type
-        else:
-            self.final_output_type = None
-
+        self.final_output = getattr(stage_config, "final_output", False)
+        self.final_output_type = getattr(stage_config, "final_output_type", None)
+        default_sampling_params = getattr(stage_config, "default_sampling_params", {})
+        self.default_sampling_params = SamplingParams(**_to_dict(default_sampling_params))
         # Runtime orchestration state (added)
         self._in_q: Optional[mp.Queue] = None
         self._out_q: Optional[mp.Queue] = None
@@ -75,6 +82,22 @@ class OmniStage:
         """Initialize the async engine for the stage."""
         self.async_engine = async_engine
 
+    def set_vllm_config(self, vllm_config) -> None:
+        """Set the vllm_config for the stage (received from worker process)."""
+        self.vllm_config = vllm_config
+
+    def set_tokenizer(self, tokenizer: AnyTokenizer) -> None:
+        """Set the tokenizer for the stage (received from worker process)."""
+        self.tokenizer = tokenizer
+
+    def set_input_preprocessor(self, input_preprocessor: InputPreprocessor) -> None:
+        """Set the input preprocessor for the stage (received from worker process)."""
+        self.input_preprocessor = input_preprocessor
+
+    def set_is_tracing_enabled(self, is_tracing_enabled: bool) -> None:
+        """Set the is_tracing_enabled for the stage (received from worker process)."""
+        self.is_tracing_enabled = is_tracing_enabled
+
     def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None:
         """Set the engine output for the stage."""
         self.engine_outputs = engine_outputs
@@ -88,6 +111,7 @@ class OmniStage:
         self,
         model: str,
         *,
+        is_async: bool = False,
         log_file: Optional[str] = None,
         shm_threshold_bytes: int = 65536,
         ctx: Optional[mp.context.BaseContext] = None,
@@ -100,16 +124,37 @@ class OmniStage:
         # Prepare lightweight dict config for worker
         engine_args = _to_dict(self.engine_args)
         runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {}))
-        stage_payload: Dict[str, Any] = {
+        stage_payload: dict[str, Any] = {
             "stage_id": self.stage_id,
             "engine_args": engine_args,
             "runtime": runtime_cfg,
             "shm_threshold_bytes": self._shm_threshold_bytes,
         }
-        self._proc = ctx.Process(
-            target=_stage_worker,
-            args=(model, stage_payload, self._in_q, self._out_q, self._log_file, batch_timeout),
-        )
+        if is_async:
+            self._proc = ctx.Process(
+                target=_stage_worker_async_entry,
+                args=(
+                    self,
+                    model,
+                    stage_payload,
+                    self._in_q,
+                    self._out_q,
+                    self._log_file,
+                    batch_timeout,
+                ),
+            )
+        else:
+            self._proc = ctx.Process(
+                target=_stage_worker,
+                args=(
+                    model,
+                    stage_payload,
+                    self._in_q,
+                    self._out_q,
+                    self._log_file,
+                    batch_timeout,
+                ),
+            )
         self._proc.start()
 
     def stop_stage_worker(self) -> None:
@@ -129,11 +174,11 @@ class OmniStage:
                 except Exception as e:
                     self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
 
-    def submit(self, payload: Dict[str, Any]) -> None:
+    def submit(self, payload: dict[str, Any]) -> None:
         assert self._in_q is not None
         self._in_q.put(payload)
 
-    def try_collect(self) -> Optional[Dict[str, Any]]:
+    def try_collect(self) -> Optional[dict[str, Any]]:
         assert self._out_q is not None
         try:
             return self._out_q.get_nowait()
@@ -142,7 +187,7 @@ class OmniStage:
 
     def process_engine_inputs(
         self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None
-    ) -> List[Union[OmniTokensPrompt, TextPrompt]]:
+    ) -> list[Union[OmniTokensPrompt, TextPrompt]]:
         """Process the engine input for the stage."""
         if self.custom_process_input_func is None:
             engine_inputs = []
@@ -150,33 +195,29 @@ class OmniStage:
                 raise ValueError("engine_input_source is empty")
             source_stage_id = self.engine_input_source[0]
             source_outputs = stage_list[source_stage_id].engine_outputs
+            if not isinstance(prompt, list):
+                prompt = [prompt]
             multi_modal_data = {
-                source_output.request_id: prompt.get("multi_modal_data", None)
-                for source_output, prompt in zip(source_outputs, prompt)
+                source_output.request_id: p.get("multi_modal_data", None)
+                for source_output, p in zip(source_outputs, prompt)
             }
 
             for source_output in source_outputs:
                 engine_input = OmniTokensPrompt(
                     prompt_token_ids=source_output.outputs[0].token_ids,
-                    multi_modal_data=(
-                        multi_modal_data[source_output.request_id]
-                        if multi_modal_data
-                        else None
-                    ),
+                    multi_modal_data=(multi_modal_data[source_output.request_id] if multi_modal_data else None),
                 )
                 engine_inputs.append(engine_input)
             return engine_inputs
 
         else:
             engine_input_source = self.engine_input_source
-            return self.custom_process_input_func(
-                stage_list, engine_input_source, prompt
-            )
+            return self.custom_process_input_func(stage_list, engine_input_source, prompt)
 
 
 def _stage_worker(
     model: str,
-    stage_payload: Dict[str, Any],
+    stage_payload: dict[str, Any],
     in_q: mp.Queue,
     out_q: mp.Queue,
     log_file: Optional[str] = None,
@@ -184,15 +225,16 @@ def _stage_worker(
 ) -> None:
     """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
     import logging as _logging
-    from vllm_omni.entrypoints.omni_llm import OmniStageLLM  # noqa: WPS433
+    import time as _time
+
     from vllm_omni.entrypoints.log_utils import (  # noqa: WPS433
-        log_stage_running_avg,
-        log_stage_batch_stats,
         compute_and_log_stage_request_stats,
         count_tokens_from_outputs,
+        log_stage_batch_stats,
+        log_stage_running_avg,
     )
-    import queue as _queue
-    import time as _time
+    from vllm_omni.entrypoints.omni_llm import OmniStageLLM  # noqa: WPS433
+
     # no inline JSONL/serialization imports; logging handled by utilities
 
     stage_id = stage_payload["stage_id"]
@@ -207,11 +249,15 @@ def _stage_worker(
             stage_log.setLevel(_logging.DEBUG)
             fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
             fh.setLevel(_logging.DEBUG)
-            fh.setFormatter(_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s"))
+            fh.setFormatter(
+                _logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")
+            )
+
             class _StageFilter(_logging.Filter):
                 def filter(self, record: _logging.LogRecord) -> bool:
                     setattr(record, "stage", stage_id)
                     return True
+
             fh.addFilter(_StageFilter())
             stage_log.addHandler(fh)
     except Exception:
@@ -233,7 +279,11 @@ def _stage_worker(
         _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
 
     # Init LLM
-    _logging.getLogger(__name__).debug("[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys()))
+    _logging.getLogger(__name__).debug(
+        "[Stage-%s] Initializing engine with args keys=%s",
+        stage_id,
+        list(engine_args.keys()),
+    )
     stage_engine = OmniStageLLM(model=model, **engine_args)
     _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
     # Signal readiness to orchestrator
@@ -251,7 +301,7 @@ def _stage_worker(
             break
 
         max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
-        batch_tasks: List[Dict[str, Any]] = [task]
+        batch_tasks: list[dict[str, Any]] = [task]
         if max_batch_size > 1:
             while len(batch_tasks) < max_batch_size:
                 if not in_q.empty():
@@ -263,11 +313,11 @@ def _stage_worker(
                 else:
                     break
 
-        batch_request_ids: List[Any] = []
-        batch_engine_inputs: List[Any] = []
-        _rx_bytes_by_rid: Dict[Any, int] = {}
-        _rx_decode_ms_by_rid: Dict[Any, float] = {}
-        _in_flight_ms_by_rid: Dict[Any, float] = {}
+        batch_request_ids: list[Any] = []
+        batch_engine_inputs: list[Any] = []
+        _rx_bytes_by_rid: dict[Any, int] = {}
+        _rx_decode_ms_by_rid: dict[Any, float] = {}
+        _in_flight_ms_by_rid: dict[Any, float] = {}
         for t in batch_tasks:
             rid = t["request_id"]
             try:
@@ -278,9 +328,7 @@ def _stage_worker(
                     _in_flight_ms_by_rid[rid] = 0.0
             except Exception:
                 _in_flight_ms_by_rid[rid] = 0.0
-            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(
-                t, obj_key="engine_inputs", shm_key="engine_inputs_shm"
-            )
+            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(t, obj_key="engine_inputs", shm_key="engine_inputs_shm")
             _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
             _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
             batch_request_ids.append(rid)
@@ -291,13 +339,21 @@ def _stage_worker(
             else:
                 _logging.getLogger(__name__).exception("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein))
         sampling_params = batch_tasks[0]["sampling_params"]
-        _logging.getLogger(__name__).debug("[Stage-%s] Received batch size=%d, request_ids=%s", stage_id, len(batch_tasks), batch_request_ids)
+        _logging.getLogger(__name__).debug(
+            "[Stage-%s] Received batch size=%d, request_ids=%s",
+            stage_id,
+            len(batch_tasks),
+            batch_request_ids,
+        )
         print("--------------------------------", flush=True)
-        print(f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}", flush=True)
+        print(
+            f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}",
+            flush=True,
+        )
         print("--------------------------------", flush=True)
         try:
             _batch_seq += 1
-            gen_outputs: List[Any] = []
+            gen_outputs: list[Any] = []
             _gen_t0 = _time.time()
             for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False):
                 gen_outputs.append(ro)
@@ -305,8 +361,8 @@ def _stage_worker(
             _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
 
             # Group outputs per request id with fallback
-            req_to_outputs: Dict[Any, List[Any]] = {rid: [] for rid in batch_request_ids}
-            unmapped: List[Any] = []
+            req_to_outputs: dict[Any, list[Any]] = {rid: [] for rid in batch_request_ids}
+            unmapped: list[Any] = []
             for ro in gen_outputs:
                 rid = getattr(ro, "request_id", None)
                 if rid in req_to_outputs:
@@ -326,12 +382,25 @@ def _stage_worker(
                 _num_tokens = count_tokens_from_outputs(_r_outputs)
                 _agg_total_tokens += _num_tokens
                 _agg_total_gen_time_ms += _gen_ms
-                _tokens_per_s = (_num_tokens * 1000.0 / _gen_ms) if _gen_ms > 0 else 0.0
 
             if _stats_file:
-                _avg_tokens_per_s = (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
-                log_stage_running_avg(_stats_file, stage_id, int(_agg_total_tokens), float(_agg_total_gen_time_ms), float(_avg_tokens_per_s))
-                log_stage_batch_stats(_stats_file, stage_id, len(batch_tasks), float(_gen_ms), list(batch_request_ids))
+                _avg_tokens_per_s = (
+                    (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
+                )
+                log_stage_running_avg(
+                    _stats_file,
+                    stage_id,
+                    int(_agg_total_tokens),
+                    float(_agg_total_gen_time_ms),
+                    float(_avg_tokens_per_s),
+                )
+                log_stage_batch_stats(
+                    _stats_file,
+                    stage_id,
+                    len(batch_tasks),
+                    float(_gen_ms),
+                    list(batch_request_ids),
+                )
 
             # Emit per-request results
             for rid in batch_request_ids:
@@ -354,25 +423,281 @@ def _stage_worker(
                             len(batch_tasks),
                             r_outputs,
                             float(_gen_ms),
-                            int(_metrics["rx_transfer_bytes"]),   # type: ignore[index]
-                            float(_metrics["rx_decode_time_ms"]), # type: ignore[index]
+                            int(_metrics["rx_transfer_bytes"]),  # type: ignore[index]
+                            float(_metrics["rx_decode_time_ms"]),  # type: ignore[index]
                         )
                     if use_shm:
-                        out_q.put({
+                        out_q.put(
+                            {
+                                "request_id": rid,
+                                "stage_id": stage_id,
+                                "engine_outputs_shm": payload,
+                                "metrics": _metrics,
+                            }
+                        )
+                    else:
+                        out_q.put(
+                            {
+                                "request_id": rid,
+                                "stage_id": stage_id,
+                                "engine_outputs": payload,
+                                "metrics": _metrics,
+                            }
+                        )
+                except Exception:
+                    out_q.put(
+                        {
+                            "request_id": rid,
+                            "stage_id": stage_id,
+                            "engine_outputs": r_outputs,
+                            "metrics": {
+                                "num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
+                                "stage_gen_time_ms": _gen_ms,
+                                "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
+                                "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
+                                "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
+                            },
+                        }
+                    )
+                _logging.getLogger(__name__).debug(
+                    "[Stage-%s] Enqueued result for request %s to downstream",
+                    stage_id,
+                    rid,
+                )
+        except Exception as e:
+            _logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e)
+            for rid in batch_request_ids:
+                out_q.put(
+                    {
+                        "request_id": rid,
+                        "stage_id": stage_id,
+                        "error": str(e),
+                    }
+                )
+
+
+def _stage_worker_async_entry(
+    omni_stage: OmniStage,
+    model: str,
+    stage_payload: dict[str, Any],
+    in_q: mp.Queue,
+    out_q: mp.Queue,
+    log_file: Optional[str] = None,
+    batch_timeout: int = 10,
+) -> None:
+    asyncio.run(_stage_worker_async(omni_stage, model, stage_payload, in_q, out_q, log_file, batch_timeout))
+
+
+async def _stage_worker_async(
+    omni_stage: OmniStage,
+    model: str,
+    stage_payload: dict[str, Any],
+    in_q: mp.Queue,
+    out_q: mp.Queue,
+    log_file: Optional[str] = None,
+    batch_timeout: int = 10,
+) -> None:
+    """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
+    import logging as _logging
+    import time as _time
+
+    from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM  # noqa: WPS433
+    from vllm_omni.entrypoints.log_utils import (  # noqa: WPS433
+        compute_and_log_stage_request_stats,
+        count_tokens_from_outputs,
+        log_stage_batch_stats,
+        log_stage_running_avg,
+    )
+
+    # no inline JSONL/serialization imports; logging handled by utilities
+
+    stage_id = stage_payload["stage_id"]
+    engine_args = stage_payload.get("engine_args", {})
+    runtime_cfg = stage_payload.get("runtime", {})
+    shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
+
+    # Per-stage file logger (optional)
+    try:
+        if log_file:
+            stage_log = _logging.getLogger(__name__)
+            stage_log.setLevel(_logging.DEBUG)
+            fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
+            fh.setLevel(_logging.DEBUG)
+            fh.setFormatter(
+                _logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")
+            )  # noqa: E501
+
+            class _StageFilter(_logging.Filter):
+                def filter(self, record: _logging.LogRecord) -> bool:
+                    setattr(record, "stage", stage_id)
+                    return True
+
+            fh.addFilter(_StageFilter())
+            stage_log.addHandler(fh)
+    except Exception:
+        pass
+
+    # Stage stats JSONL file
+    _stats_file = f"{log_file}.stage{stage_id}.stats.jsonl" if log_file else None
+
+    # Aggregates for running average
+    _agg_total_tokens = 0
+    _agg_total_gen_time_ms = 0.0
+    # Monotonic batch id per stage process for orchestrator dedup on time
+    # aggregation
+    _batch_seq = 0
+
+    # Device mapping
+    try:
+        set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
+    except Exception as e:
+        _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
+
+    # Init LLM
+    _logging.getLogger(__name__).debug(
+        "[Stage-%s] Initializing engine with args keys=%s",
+        stage_id,
+        list(engine_args.keys()),
+    )
+    omni_engine_args = AsyncOmniEngineArgs(model=model, **engine_args)
+    usage_context = UsageContext.OPENAI_API_SERVER
+    vllm_config = omni_engine_args.create_engine_config(usage_context=usage_context)
+    stage_engine = AsyncOmniStageLLM.from_vllm_config(
+        vllm_config=vllm_config,
+        usage_context=usage_context,
+        engine_args=omni_engine_args,
+    )
+    omni_stage.set_async_engine(stage_engine)
+    # Don't keep the dummy data in memory
+    await stage_engine.reset_mm_cache()
+    _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
+    # Signal readiness to orchestrator and send vllm_config back to main process
+    try:
+        # Send vllm_config back to main process so it can be accessed via
+        # get_vllm_config(). This is needed because async_engine is only available
+        # in the worker process
+
+        # input_preprocessor = await stage_engine.get_input_preprocessor()
+        out_q.put(
+            {
+                "type": "stage_ready",
+                "stage_id": stage_id,
+                "vllm_config": vllm_config,
+                "tokenizer": getattr(stage_engine, "tokenizer", None),
+                "is_tracing_enabled": await stage_engine.is_tracing_enabled(),
+                # "input_preprocessor": input_preprocessor,
+            }
+        )
+    except Exception as e:
+        _logging.getLogger(__name__).warning("[Stage-%s] Failed to send stage ready signal: %s", stage_id, e)
+
+    # Batch processing loop
+    while True:
+        task = in_q.get()
+        _recv_dequeue_ts = _time.time()
+        if task is None:
+            _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id)
+            break
+
+        _rx_bytes_by_rid: dict[Any, int] = {}
+        _rx_decode_ms_by_rid: dict[Any, float] = {}
+        _in_flight_ms_by_rid: dict[Any, float] = {}
+
+        rid = task["request_id"]
+        try:
+            sent_ts = float(task.get("sent_ts", None)) if isinstance(task, dict) else None
+            if sent_ts is not None:
+                _in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0
+            else:
+                _in_flight_ms_by_rid[rid] = 0.0
+        except Exception:
+            _in_flight_ms_by_rid[rid] = 0.0
+        ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+        _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
+        _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
+
+        sampling_params = task["sampling_params"]
+        _logging.getLogger(__name__).debug("[Stage-%s] Received batch size=1, request_ids=%d", stage_id, rid)
+        print("--------------------------------", flush=True)
+        print(f"[Stage-{stage_id}] Received batch size=1, request_ids={rid}", flush=True)
+        print("--------------------------------", flush=True)
+        try:
+            _batch_seq += 1
+            _gen_t0 = _time.time()
+            if isinstance(ein, list):
+                ein = ein[0]
+
+            async for res in stage_engine.generate(ein, sampling_params, rid):
+                gen_output = res
+            _gen_t1 = _time.time()
+            _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
+
+            r_outputs = [gen_output]
+            _num_tokens = count_tokens_from_outputs(r_outputs)
+            _agg_total_tokens += _num_tokens
+            _agg_total_gen_time_ms += _gen_ms
+
+            if _stats_file:
+                _avg_tokens_per_s = (
+                    (_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
+                )
+                log_stage_running_avg(
+                    _stats_file,
+                    stage_id,
+                    int(_agg_total_tokens),
+                    float(_agg_total_gen_time_ms),
+                    float(_avg_tokens_per_s),
+                )
+                log_stage_batch_stats(_stats_file, stage_id, 1, float(_gen_ms), [rid])
+
+            try:
+                use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
+                _metrics = {
+                    "num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
+                    "stage_gen_time_ms": _gen_ms,
+                    "batch_id": int(_batch_seq),
+                    "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
+                    "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
+                    "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
+                }
+                if _stats_file:
+                    compute_and_log_stage_request_stats(
+                        _stats_file,
+                        stage_id,
+                        rid,
+                        1,
+                        r_outputs,
+                        float(_gen_ms),
+                        int(_metrics["rx_transfer_bytes"]),  # type: ignore[index]
+                        float(_metrics["rx_decode_time_ms"]),  # type: ignore[index]
+                    )
+                if use_shm:
+                    out_q.put(
+                        {
                             "request_id": rid,
                             "stage_id": stage_id,
                             "engine_outputs_shm": payload,
                             "metrics": _metrics,
-                        })
-                    else:
-                        out_q.put({
+                        }
+                    )
+                else:
+                    out_q.put(
+                        {
                             "request_id": rid,
                             "stage_id": stage_id,
                             "engine_outputs": payload,
                             "metrics": _metrics,
-                        })
-                except Exception:
-                    out_q.put({
+                        }
+                    )
+            except Exception as e:
+                _logging.getLogger(__name__).exception(
+                    "[Stage-%s] Failed to enqueue result for request %s: %s",
+                    stage_id,
+                    rid,
+                    e,
+                )
+                out_q.put(
+                    {
                         "request_id": rid,
                         "stage_id": stage_id,
                         "engine_outputs": r_outputs,
@@ -383,13 +708,16 @@ def _stage_worker(
                             "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
                             "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
                         },
-                    })
-                _logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid)
+                    }
+                )
+            _logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid)
+
         except Exception as e:
-            _logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e)
-            for rid in batch_request_ids:
-                out_q.put({
+            _logging.getLogger(__name__).exception("[Stage-%s] Failed on request %s: %s", stage_id, rid, e)
+            out_q.put(
+                {
                     "request_id": rid,
                     "stage_id": stage_id,
                     "error": str(e),
-                })
+                }
+            )

vllm_omni/entrypoints/openai/api_server.py

@@ -0,0 +1,313 @@
+import multiprocessing
+import multiprocessing.forkserver as forkserver
+import os
+from argparse import Namespace
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from http import HTTPStatus
+from typing import Any, Optional
+
+import vllm.envs as envs
+from fastapi import Depends, HTTPException, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+from starlette.datastructures import State
+from vllm.config import VllmConfig
+from vllm.engine.protocol import EngineClient
+from vllm.entrypoints.chat_utils import load_chat_template, resolve_hf_chat_template, resolve_mistral_chat_template
+from vllm.entrypoints.launcher import serve_http
+from vllm.entrypoints.logger import RequestLogger
+from vllm.entrypoints.openai.api_server import (
+    base,
+    build_app,
+    load_log_config,
+    maybe_register_tokenizer_info_endpoint,
+    router,
+    setup_server,
+    validate_json_request,
+)
+from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ChatCompletionResponse, ErrorResponse
+from vllm.entrypoints.openai.serving_models import BaseModelPath, LoRAModulePath, OpenAIServingModels
+from vllm.entrypoints.openai.tool_parsers import ToolParserManager
+
+# yapf conflicts with isort for this block
+# yapf: disable
+# yapf: enable
+from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer
+from vllm.entrypoints.utils import load_aware_call, with_cancellation
+from vllm.logger import init_logger
+from vllm.transformers_utils.tokenizer import MistralTokenizer
+from vllm.utils import decorate_logs
+
+from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+logger = init_logger(__name__)
+
+
+async def omni_run_server(args, **uvicorn_kwargs) -> None:
+    """Run a single-worker API server."""
+
+    # Add process-specific prefix to stdout and stderr.
+    decorate_logs("APIServer")
+
+    listen_address, sock = setup_server(args)
+    await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
+
+
+async def omni_run_server_worker(listen_address, sock, args, client_config=None, **uvicorn_kwargs) -> None:
+    """Run a single API server worker."""
+
+    if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
+        ToolParserManager.import_tool_parser(args.tool_parser_plugin)
+
+    # Load logging config for uvicorn if specified
+    log_config = load_log_config(args.log_config_file)
+    if log_config is not None:
+        uvicorn_kwargs["log_config"] = log_config
+
+    async with build_async_omni_llm(
+        args,
+        client_config=client_config,
+    ) as engine_client:
+        maybe_register_tokenizer_info_endpoint(args)
+        app = build_app(args)
+
+        vllm_config = await engine_client.get_vllm_config()
+        await omni_init_app_state(engine_client, vllm_config, app.state, args)
+
+        logger.info(
+            "Starting vLLM API server %d on %s",
+            vllm_config.parallel_config._api_process_rank,
+            listen_address,
+        )
+        shutdown_task = await serve_http(
+            app,
+            sock=sock,
+            enable_ssl_refresh=args.enable_ssl_refresh,
+            host=args.host,
+            port=args.port,
+            log_level=args.uvicorn_log_level,
+            # NOTE: When the 'disable_uvicorn_access_log' value is True,
+            # no access log will be output.
+            access_log=not args.disable_uvicorn_access_log,
+            timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE,
+            ssl_keyfile=args.ssl_keyfile,
+            ssl_certfile=args.ssl_certfile,
+            ssl_ca_certs=args.ssl_ca_certs,
+            ssl_cert_reqs=args.ssl_cert_reqs,
+            h11_max_incomplete_event_size=args.h11_max_incomplete_event_size,
+            h11_max_header_count=args.h11_max_header_count,
+            **uvicorn_kwargs,
+        )
+
+    # NB: Await server shutdown only after the backend context is exited
+    try:
+        await shutdown_task
+    finally:
+        sock.close()
+
+
+@asynccontextmanager
+async def build_async_omni_llm(
+    args: Namespace,
+    *,
+    disable_frontend_multiprocessing: Optional[bool] = None,
+    client_config: Optional[dict[str, Any]] = None,
+) -> AsyncIterator[EngineClient]:
+    if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver":
+        # The executor is expected to be mp.
+        # Pre-import heavy modules in the forkserver process
+        logger.debug("Setup forkserver with pre-imports")
+        multiprocessing.set_start_method("forkserver")
+        multiprocessing.set_forkserver_preload(["vllm.v1.engine.async_llm"])
+        forkserver.ensure_running()
+        logger.debug("Forkserver setup complete!")
+
+    # Context manager to handle async_omni_llm lifecycle
+    # Ensures everything is shutdown and cleaned up on error/exit
+    async with build_async_omni_llm_from_stage_config(
+        args,
+        disable_frontend_multiprocessing=disable_frontend_multiprocessing,
+    ) as async_omni_llm:
+        yield async_omni_llm
+
+
+@asynccontextmanager
+async def build_async_omni_llm_from_stage_config(
+    args: Namespace,
+    *,
+    disable_frontend_multiprocessing: bool = False,
+    client_config: Optional[dict[str, Any]] = None,
+) -> AsyncIterator[EngineClient]:
+    """
+    Create AsyncOmniLLM, either:
+        - in-process using the AsyncOmniLLM Directly
+        - multiprocess using AsyncOmniLLM RPC
+
+    Returns the AsyncOmniLLM or None if the creation failed.
+    """
+
+    # V1 AsyncLLM.
+    assert envs.VLLM_USE_V1
+
+    if disable_frontend_multiprocessing:
+        logger.warning(
+            "V1 is enabled, but got --disable-frontend-multiprocessing. "
+            "To disable frontend multiprocessing, set VLLM_USE_V1=0."
+        )
+
+    async_omni_llm: Optional[EngineClient] = None
+
+    try:
+        if getattr(args, "stage_configs_path", None):
+            async_omni_llm = AsyncOmniLLM(model=args.model, stage_configs_path=args.stage_configs_path)
+        else:
+            async_omni_llm = AsyncOmniLLM(model=args.model)
+
+        # # Don't keep the dummy data in memory
+        # await async_llm.reset_mm_cache()
+
+        yield async_omni_llm
+    finally:
+        if async_omni_llm:
+            async_omni_llm.shutdown()
+
+
+async def omni_init_app_state(
+    engine_client: EngineClient,
+    vllm_config: VllmConfig,
+    state: State,
+    args: Namespace,
+) -> None:
+    if args.served_model_name is not None:
+        served_model_names = args.served_model_name
+    else:
+        served_model_names = [args.model]
+
+    if args.enable_log_requests:
+        request_logger = RequestLogger(max_log_len=args.max_log_len)
+    else:
+        request_logger = None
+
+    base_model_paths = [BaseModelPath(name=name, model_path=args.model) for name in served_model_names]
+    state.engine_client = engine_client
+    state.log_stats = not args.disable_log_stats
+    state.vllm_config = vllm_config
+    model_config = vllm_config.model_config
+    state.log_stats = not args.disable_log_stats
+
+    # For omni models
+    state.stage_configs = engine_client.stage_configs
+
+    resolved_chat_template = load_chat_template(args.chat_template)
+    if resolved_chat_template is not None:
+        # Get the tokenizer to check official template
+        tokenizer = await engine_client.get_tokenizer()
+
+        if isinstance(tokenizer, MistralTokenizer):
+            # The warning is logged in resolve_mistral_chat_template.
+            resolved_chat_template = resolve_mistral_chat_template(chat_template=resolved_chat_template)
+        else:
+            hf_chat_template = resolve_hf_chat_template(
+                tokenizer=tokenizer,
+                chat_template=None,
+                tools=None,
+                model_config=vllm_config.model_config,
+            )
+
+            if hf_chat_template != resolved_chat_template:
+                logger.warning(
+                    "Using supplied chat template: %s\nIt is different from official chat template '%s'. This discrepancy may lead to performance degradation.",  # noqa: E501
+                    resolved_chat_template,
+                    args.model,
+                )
+
+    if args.tool_server == "demo":
+        tool_server: Optional[ToolServer] = DemoToolServer()
+        assert isinstance(tool_server, DemoToolServer)
+        await tool_server.init_and_validate()
+    elif args.tool_server:
+        tool_server = MCPToolServer()
+        await tool_server.add_tool_server(args.tool_server)
+    else:
+        tool_server = None
+
+    # Merge default_mm_loras into the static lora_modules
+    default_mm_loras = vllm_config.lora_config.default_mm_loras if vllm_config.lora_config is not None else {}
+
+    lora_modules = args.lora_modules
+    if default_mm_loras:
+        default_mm_lora_paths = [
+            LoRAModulePath(
+                name=modality,
+                path=lora_path,
+            )
+            for modality, lora_path in default_mm_loras.items()
+        ]
+        if args.lora_modules is None:
+            lora_modules = default_mm_lora_paths
+        else:
+            lora_modules += default_mm_lora_paths
+
+    state.openai_serving_models = OpenAIServingModels(
+        engine_client=engine_client,
+        model_config=model_config,
+        base_model_paths=base_model_paths,
+        lora_modules=lora_modules,
+    )
+    await state.openai_serving_models.init_static_loras()
+    state.openai_serving_chat = OmniOpenAIServingChat(
+        engine_client,
+        model_config,
+        state.openai_serving_models,
+        args.response_role,
+        request_logger=request_logger,
+        chat_template=resolved_chat_template,
+        chat_template_content_format=args.chat_template_content_format,
+        trust_request_chat_template=args.trust_request_chat_template,
+        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
+        enable_auto_tools=args.enable_auto_tool_choice,
+        exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
+        tool_parser=args.tool_call_parser,
+        reasoning_parser=args.structured_outputs_config.reasoning_parser,
+        enable_prompt_tokens_details=args.enable_prompt_tokens_details,
+        enable_force_include_usage=args.enable_force_include_usage,
+        enable_log_outputs=args.enable_log_outputs,
+        log_error_stack=args.log_error_stack,
+    )
+
+    state.enable_server_load_tracking = args.enable_server_load_tracking
+    state.server_load_metrics = 0
+
+
+def Omnichat(request: Request) -> Optional[OmniOpenAIServingChat]:
+    return request.app.state.openai_serving_chat
+
+
+@router.post(
+    "/v1/chat/completions",
+    dependencies=[Depends(validate_json_request)],
+    responses={
+        HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+@with_cancellation
+@load_aware_call
+async def create_chat_completion(request: ChatCompletionRequest, raw_request: Request):
+    handler = Omnichat(raw_request)
+    if handler is None:
+        return base(raw_request).create_error_response(message="The model does not support Chat Completions API")
+    try:
+        generator = await handler.create_chat_completion(request, raw_request)
+    except Exception as e:
+        raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
+    if isinstance(generator, ErrorResponse):
+        return JSONResponse(content=generator.model_dump(), status_code=generator.error.code)
+
+    elif isinstance(generator, ChatCompletionResponse):
+        return JSONResponse(content=generator.model_dump())
+
+    return StreamingResponse(content=generator, media_type="text/event-stream")

vllm_omni/entrypoints/openai/serving_chat.py

@@ -0,0 +1,643 @@
+import asyncio
+import base64
+import json
+import time
+import uuid
+from collections.abc import AsyncGenerator, AsyncIterator
+from datetime import datetime, timedelta, timezone
+from io import BytesIO
+from typing import Optional, Union
+
+import jinja2
+from fastapi import Request
+from pydantic import TypeAdapter
+
+try:
+    import soundfile
+except ImportError:
+    soundfile = None
+
+from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
+from vllm.entrypoints.chat_utils import ConversationMessage, get_history_tool_calls_cnt, make_tool_call_id
+from vllm.entrypoints.harmony_utils import parse_chat_output
+from vllm.entrypoints.openai.protocol import (
+    ChatCompletionNamedToolChoiceParam,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice,
+    ChatMessage,
+    ErrorResponse,
+    FunctionCall,
+    FunctionDefinition,
+    PromptTokenUsageInfo,
+    RequestResponseMetadata,
+    ToolCall,
+    UsageInfo,
+)
+from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
+from vllm.entrypoints.openai.serving_engine import RequestPrompt, clamp_prompt_logprobs
+from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import MistralToolCall
+from vllm.inputs.data import PromptType
+from vllm.logger import init_logger
+from vllm.lora.request import LoRARequest
+from vllm.outputs import RequestOutput
+from vllm.sampling_params import SamplingParams
+from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
+from vllm.transformers_utils.tokenizers import (
+    maybe_serialize_tool_calls,
+    truncate_tool_call_ids,
+    validate_request_params,
+)
+from vllm.utils import as_list
+
+from vllm_omni.outputs import OmniRequestOutput
+
+logger = init_logger(__name__)
+
+
+class OmniOpenAIServingChat(OpenAIServingChat):
+    async def create_chat_completion(
+        self,
+        request: ChatCompletionRequest,
+        raw_request: Optional[Request] = None,
+    ) -> Union[AsyncGenerator[str, None], ChatCompletionResponse, ErrorResponse]:
+        """
+        Chat Completion API similar to OpenAI's API.
+
+        See https://platform.openai.com/docs/api-reference/chat/create
+        for the API specification. This API mimics the OpenAI
+        Chat Completion API.
+        """
+        error_check_ret = await self._check_model(request)
+        if error_check_ret is not None:
+            logger.error("Error with model %s", error_check_ret)
+            return error_check_ret
+
+        # If the engine is dead, raise the engine's DEAD_ERROR.
+        # This is required for the streaming case, where we return a
+        # success status before we actually start generating text :).
+        if self.engine_client.errored:
+            raise self.engine_client.dead_error
+
+        try:
+            lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True)
+
+            model_name = self.models.model_name(lora_request)
+
+            tokenizer = await self.engine_client.get_tokenizer()
+
+            tool_parser = self.tool_parser
+
+            if isinstance(tokenizer, MistralTokenizer):
+                # because of issues with pydantic we need to potentially
+                # re-serialize the tool_calls field of the request
+                # for more info: see comment in `maybe_serialize_tool_calls`
+                maybe_serialize_tool_calls(request)
+                truncate_tool_call_ids(request)
+                validate_request_params(request)
+
+            if (
+                request.tool_choice == "auto"
+                and not (self.enable_auto_tools and tool_parser is not None)
+                and not isinstance(tokenizer, MistralTokenizer)
+                and not self.use_harmony
+            ):
+                # for hf tokenizers, "auto" tools requires
+                # --enable-auto-tool-choice and --tool-call-parser
+                return self.create_error_response(
+                    '"auto" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set'
+                )
+
+            if request.tools is None or (request.tool_choice == "none" and self.exclude_tools_when_tool_choice_none):
+                tool_dicts = None
+            else:
+                tool_dicts = [tool.model_dump() for tool in request.tools]
+
+            # Common case.
+            request_chat_template = request.chat_template
+            chat_template_kwargs = request.chat_template_kwargs
+            if not self.trust_request_chat_template and (
+                request_chat_template is not None
+                or (chat_template_kwargs and chat_template_kwargs.get("chat_template") is not None)
+            ):
+                return self.create_error_response(
+                    "Chat template is passed with request, but --trust-request-chat-template is not set. "
+                    "Refused request with untrusted chat template."
+                )
+            (
+                conversation,
+                request_prompts,
+                engine_prompts,
+            ) = await self._preprocess_chat(
+                request,
+                tokenizer,
+                request.messages,
+                chat_template=request_chat_template or self.chat_template,
+                chat_template_content_format=self.chat_template_content_format,
+                add_generation_prompt=request.add_generation_prompt,
+                continue_final_message=request.continue_final_message,
+                tool_dicts=tool_dicts,
+                documents=request.documents,
+                chat_template_kwargs=request.chat_template_kwargs,
+                tool_parser=tool_parser,
+                add_special_tokens=request.add_special_tokens,
+            )
+
+        except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e:
+            logger.exception("Error in preprocessing prompt inputs")
+            return self.create_error_response(f"{e} {e.__cause__}")
+
+        request_id = f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}"
+
+        request_metadata = RequestResponseMetadata(request_id=request_id)
+        if raw_request:
+            raw_request.state.request_metadata = request_metadata
+
+        # Schedule the request and get the result generator.
+        generators: list[AsyncGenerator[RequestOutput, None]] = []
+        try:
+            for i, engine_prompt in enumerate(engine_prompts):
+                if hasattr(request, "sampling_params_list"):
+                    sampling_params_list = self._to_sampling_params_list(request.sampling_params_list)
+                else:
+                    sampling_params_list = None
+
+                self._log_inputs(
+                    request_id,
+                    request_prompts[i],
+                    params_list=sampling_params_list,
+                    lora_request=lora_request,
+                )
+
+                trace_headers = None if raw_request is None else await self._get_trace_headers(raw_request.headers)
+
+                generator = self.engine_client.generate(
+                    prompt=engine_prompt,
+                    request_id=request_id,
+                    sampling_params_list=sampling_params_list,
+                    lora_request=lora_request,
+                    trace_headers=trace_headers,
+                    priority=request.priority,
+                )
+
+                generators.append(generator)
+        except ValueError as e:
+            # TODO: Use a vllm-specific Validation Error
+            return self.create_error_response(str(e))
+
+        assert len(generators) == 1
+        (result_generator,) = generators
+
+        # Streaming response
+        if request.stream:
+            return self.chat_completion_stream_generator(
+                request,
+                result_generator,
+                request_id,
+                model_name,
+                conversation,
+                tokenizer,
+                request_metadata,
+                enable_force_include_usage=self.enable_force_include_usage,
+            )
+
+        try:
+            return await self.chat_completion_full_generator(
+                request,
+                result_generator,
+                request_id,
+                model_name,
+                conversation,
+                tokenizer,
+                request_metadata,
+            )
+        except ValueError as e:
+            # TODO: Use a vllm-specific Validation Error
+            return self.create_error_response(str(e))
+
+    def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]:
+        final_sampling_params_list = []
+        for sampling_params in sampling_params_list:
+            if isinstance(sampling_params, dict):
+                final_sampling_params_list.append(SamplingParams(**sampling_params))
+            elif isinstance(sampling_params, SamplingParams):
+                final_sampling_params_list.append(sampling_params)
+            else:
+                raise ValueError(f"Invalid sampling params: {sampling_params}")
+        return final_sampling_params_list
+
+    def _log_inputs(
+        self,
+        request_id: str,
+        inputs: Union[RequestPrompt, PromptType],
+        params_list: Optional[list[SamplingParams]],
+        lora_request: Optional[LoRARequest],
+    ) -> None:
+        if self.request_logger is None:
+            return
+        prompt, prompt_token_ids, prompt_embeds = None, None, None
+        if isinstance(inputs, str):
+            prompt = inputs
+        elif isinstance(inputs, list):
+            prompt_token_ids = inputs
+        else:
+            prompt = getattr(inputs, "prompt", None)
+            prompt_token_ids = getattr(inputs, "prompt_token_ids", None)
+
+        logger.info(
+            "Received request %s: prompt: %r, params_list: %s, prompt_token_ids: %s, prompt_embeds shape: %s, lora_request: %s.",  # noqa: E501
+            request_id,
+            prompt,
+            params_list,
+            prompt_token_ids,
+            prompt_embeds.shape if prompt_embeds is not None else None,
+            lora_request,
+        )
+
+    async def chat_completion_full_generator(
+        self,
+        request: ChatCompletionRequest,
+        result_generator: AsyncIterator[RequestOutput],
+        request_id: str,
+        model_name: str,
+        conversation: list[ConversationMessage],
+        tokenizer: AnyTokenizer,
+        request_metadata: RequestResponseMetadata,
+    ) -> Union[ErrorResponse, ChatCompletionResponse]:
+        created_time = int(time.time())
+        final_res: Optional[RequestOutput] = None
+
+        final_outputs: list[OmniRequestOutput] = []
+        try:
+            async for res in result_generator:
+                final_outputs.append(res)
+        except asyncio.CancelledError:
+            return self.create_error_response("Client disconnected")
+        except ValueError as e:
+            # TODO: Use a vllm-specific Validation Error
+            return self.create_error_response(str(e))
+
+        assert final_outputs is not None
+
+        choices: list[ChatCompletionResponseChoice] = []
+
+        usage = UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0)
+        role = self.get_chat_request_role(request)
+        prompt_logprobs = None
+        prompt_token_ids = None
+        kv_transfer_params = None
+
+        for omni_outputs in final_outputs:
+            choices_data = []
+            if omni_outputs.final_output_type == "text":
+                (
+                    choices_data,
+                    usage,
+                    prompt_logprobs,
+                    prompt_token_ids,
+                    kv_transfer_params,
+                ) = self._create_text_choice(request, omni_outputs, tokenizer, conversation, role)
+            elif omni_outputs.final_output_type == "audio":
+                choices_data = self._create_audio_choice(omni_outputs, role)
+            else:
+                logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}")
+                continue
+            choices.extend(choices_data)
+
+        response = ChatCompletionResponse(
+            id=request_id,
+            created=created_time,
+            model=model_name,
+            choices=choices,
+            usage=usage,
+            prompt_logprobs=prompt_logprobs,
+            prompt_token_ids=prompt_token_ids,
+            kv_transfer_params=kv_transfer_params,
+        )
+
+        # Log complete response if output logging is enabled
+        if self.enable_log_outputs and self.request_logger:
+            for choice in choices:
+                output_text = ""
+                if choice.message.content:
+                    output_text = choice.message.content
+                elif choice.message.tool_calls:
+                    # For tool calls, log the function name and arguments
+                    tool_call_descriptions = []
+                    for tc in choice.message.tool_calls:
+                        if hasattr(tc.function, "name") and hasattr(tc.function, "arguments"):
+                            tool_call_descriptions.append(f"{tc.function.name}({tc.function.arguments})")
+                    tool_calls_str = ", ".join(tool_call_descriptions)
+                    output_text = f"[tool_calls: {tool_calls_str}]"
+
+                if output_text:
+                    # Get the corresponding output token IDs
+                    output_token_ids = None
+                    if choice.index < len(final_res.outputs):
+                        output_token_ids = final_res.outputs[choice.index].token_ids
+
+                    self.request_logger.log_outputs(
+                        request_id=request_id,
+                        outputs=output_text,
+                        output_token_ids=output_token_ids,
+                        finish_reason=choice.finish_reason,
+                        is_streaming=False,
+                        delta=False,
+                    )
+
+        return response
+
+    def _create_text_choice(
+        self,
+        request: ChatCompletionRequest,
+        omni_outputs: OmniRequestOutput,
+        tokenizer: AnyTokenizer,
+        conversation: list[ConversationMessage],
+        role: str,
+    ):
+        final_res = omni_outputs.request_output
+        if self.tool_call_id_type == "kimi_k2":
+            history_tool_call_cnt = get_history_tool_calls_cnt(conversation)
+        else:
+            history_tool_call_cnt = 0
+
+        choices: list[ChatCompletionResponseChoice] = []
+
+        for output in final_res.outputs:
+            token_ids = output.token_ids
+            out_logprobs = output.logprobs
+            tool_call_info = None
+
+            if request.logprobs and request.top_logprobs is not None:
+                assert out_logprobs is not None, "Did not output logprobs"
+                logprobs = self._create_chat_logprobs(
+                    token_ids=token_ids,
+                    top_logprobs=out_logprobs,
+                    num_output_top_logprobs=request.top_logprobs,
+                    tokenizer=tokenizer,
+                    return_as_token_id=request.return_tokens_as_token_ids,
+                )
+            else:
+                logprobs = None
+
+            if self.use_harmony:
+                reasoning_content, content, _ = parse_chat_output(token_ids)
+                if not request.include_reasoning:
+                    reasoning_content = None
+
+                if self.tool_parser is not None:
+                    tool_parser = self.tool_parser(tokenizer)
+                    # NOTE: We use token_ids for openai tool parser
+                    tool_call_info = tool_parser.extract_tool_calls(
+                        "",
+                        request=request,
+                        token_ids=token_ids,  # type: ignore
+                    )
+                    content = tool_call_info.content
+                    message = ChatMessage(
+                        role=role,
+                        reasoning_content=reasoning_content,
+                        content=content,
+                        tool_calls=tool_call_info.tool_calls,
+                    )
+                else:
+                    message = ChatMessage(
+                        role=role,
+                        reasoning_content=reasoning_content,
+                        content=content,
+                    )
+
+                choice_data = ChatCompletionResponseChoice(
+                    index=output.index,
+                    message=message,
+                    logprobs=logprobs,
+                    finish_reason=(
+                        "tool_calls"
+                        if (tool_call_info is not None and tool_call_info.tools_called)
+                        else (output.finish_reason if output.finish_reason else "stop")
+                    ),
+                    stop_reason=output.stop_reason,
+                )
+                choices.append(choice_data)
+                continue
+
+            if self.reasoning_parser:
+                try:
+                    reasoning_parser = self.reasoning_parser(tokenizer)
+                except RuntimeError as e:
+                    logger.exception("Error in reasoning parser creation.")
+                    return self.create_error_response(str(e))
+                # If the reasoning parser is enabled,
+                # tool calls are extracted exclusively from the content.
+                reasoning_content, content = reasoning_parser.extract_reasoning_content(output.text, request=request)
+                if not request.include_reasoning:
+                    reasoning_content = None
+            else:
+                reasoning_content = None
+                content = output.text
+
+            auto_tools_called = False
+            # if auto tools are not enabled, and a named tool choice using
+            #   outlines is not being used
+            if (not self.enable_auto_tools or not self.tool_parser) and (
+                not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
+                and request.tool_choice != "required"
+            ):
+                message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
+
+            # if the request uses tools and specified a tool choice
+            elif request.tool_choice and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam:
+                tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall
+                message = ChatMessage(
+                    role=role,
+                    reasoning_content=reasoning_content,
+                    content="",
+                    tool_calls=[
+                        tool_call_class(
+                            function=FunctionCall(
+                                name=request.tool_choice.function.name,
+                                arguments=content,
+                            )
+                        )
+                    ],
+                )
+
+            elif request.tool_choice and request.tool_choice == "required":
+                tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall
+
+                # the fields of FunctionDefinition are a superset of the
+                # tool call outputs and can be used for parsing
+                assert content is not None
+                tool_calls = TypeAdapter(list[FunctionDefinition]).validate_json(content)
+                tool_call_ids = []
+                for tool_call in tool_calls:
+                    tool_call_ids.append(
+                        make_tool_call_id(
+                            id_type=self.tool_call_id_type,
+                            func_name=tool_call.name,
+                            idx=history_tool_call_cnt,
+                        )
+                    )
+                    history_tool_call_cnt += 1
+                message = ChatMessage(
+                    role=role,
+                    content="",
+                    tool_calls=[
+                        tool_call_class(
+                            id=tool_call_ids[i],
+                            function=FunctionCall(
+                                name=tool_call.name,
+                                arguments=json.dumps(tool_call.parameters, ensure_ascii=False),
+                            ),
+                        )
+                        for i, tool_call in enumerate(tool_calls)
+                    ],
+                    reasoning_content=reasoning_content,
+                )
+
+            # if the request doesn't use tool choice
+            # OR specifies to not use a tool
+            elif not request.tool_choice or request.tool_choice == "none":
+                message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
+
+            # handle when there are tools and tool choice is auto
+            elif (
+                request.tools
+                and (request.tool_choice == "auto" or request.tool_choice is None)
+                and self.enable_auto_tools
+                and self.tool_parser
+            ):
+                try:
+                    tool_parser = self.tool_parser(tokenizer)
+                except RuntimeError as e:
+                    logger.exception("Error in tool parser creation.")
+                    return self.create_error_response(str(e))
+
+                tool_call_info = tool_parser.extract_tool_calls(content if content is not None else "", request=request)
+                # In the OpenAI API the finish_reason is "tools_called"
+                # if the tool choice is auto and the model produced a tool
+                # call. The same is not true for named function calls
+                auto_tools_called = tool_call_info.tools_called
+                if tool_call_info.tools_called:
+                    message = ChatMessage(
+                        role=role,
+                        reasoning_content=reasoning_content,
+                        content=tool_call_info.content,
+                        tool_calls=tool_call_info.tool_calls,
+                    )
+
+                else:
+                    # FOR NOW make it a chat message; we will have to detect
+                    # the type to make it later.
+                    ret_content = content
+
+                    # try to use content return from tool parser first,
+                    # tool parser may do some modify for the content.
+                    if tool_call_info.content and len(tool_call_info.content) > 0:
+                        ret_content = tool_call_info.content
+                    message = ChatMessage(
+                        role=role,
+                        reasoning_content=reasoning_content,
+                        content=ret_content,
+                    )
+
+            # undetermined case that is still important to handle
+            else:
+                logger.error(
+                    "Error in chat_completion_full_generator - cannot determine if tools should be extracted. "
+                    "Returning a standard chat completion."
+                )
+                message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
+
+            choice_data = ChatCompletionResponseChoice(
+                index=output.index,
+                message=message,
+                logprobs=logprobs,
+                finish_reason=(
+                    "tool_calls" if auto_tools_called else output.finish_reason if output.finish_reason else "stop"
+                ),
+                stop_reason=output.stop_reason,
+                token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
+            )
+            choices.append(choice_data)
+
+        if request.echo:
+            last_msg_content: Union[str, list[dict[str, str]]] = ""
+            if conversation and "content" in conversation[-1] and conversation[-1].get("role") == role:
+                last_msg_content = conversation[-1]["content"] or ""
+            if isinstance(last_msg_content, list):
+                last_msg_content = "\n".join(msg["text"] for msg in last_msg_content)
+
+            for choice in choices:
+                full_message = last_msg_content + (choice.message.content or "")
+                choice.message.content = full_message
+
+        assert final_res.prompt_token_ids is not None
+        num_prompt_tokens = len(final_res.prompt_token_ids)
+        if final_res.encoder_prompt_token_ids is not None:
+            num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
+        num_generated_tokens = sum(len(output.token_ids) for output in final_res.outputs)
+        usage = UsageInfo(
+            prompt_tokens=num_prompt_tokens,
+            completion_tokens=num_generated_tokens,
+            total_tokens=num_prompt_tokens + num_generated_tokens,
+        )
+        if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
+            usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=final_res.num_cached_tokens)
+
+        prompt_logprobs = clamp_prompt_logprobs(final_res.prompt_logprobs)
+        prompt_token_ids = final_res.prompt_token_ids if request.return_token_ids else None
+        kv_transfer_params = final_res.kv_transfer_params
+
+        return choices, usage, prompt_logprobs, prompt_token_ids, kv_transfer_params
+
+    def _create_audio_choice(self, omni_outputs: OmniRequestOutput, role: str):
+        choices: list[ChatCompletionResponseChoice] = []
+        final_res = omni_outputs.request_output
+        audio_tensor = final_res.multimodal_output["audio"].detach().cpu().numpy()
+
+        # Convert numpy array to WAV bytes and encode as base64
+        if soundfile is None:
+            raise ImportError(
+                "soundfile is required for audio generation. Please install it with: pip install soundfile"
+            )
+
+        # Default sample rate for TTS models (typically 24000 Hz)
+        # You may need to adjust this based on your model's configuration
+        sample_rate = 24000
+
+        # Ensure audio is 1D (flatten if needed)
+        if audio_tensor.ndim > 1:
+            audio_tensor = audio_tensor.flatten()
+
+        # Convert to WAV format and encode as base64
+        with BytesIO() as buffer:
+            soundfile.write(buffer, audio_tensor, sample_rate, format="WAV")
+            wav_bytes = buffer.getvalue()
+
+        audio_base64 = base64.b64encode(wav_bytes).decode("utf-8")
+
+        # Generate unique ID for the audio
+        audio_id = f"audio-{uuid.uuid4().hex[:16]}"
+
+        # Set expiration time (e.g., 24 hours from now) as Unix timestamp
+        expires_at = int((datetime.now(timezone.utc) + timedelta(hours=24)).timestamp())
+
+        # Create OpenAIChatCompletionAudio object with all required fields
+        audio_obj = OpenAIChatCompletionAudio(
+            id=audio_id,
+            data=audio_base64,
+            expires_at=expires_at,
+            transcript="",  # Empty transcript if not available
+        )
+
+        for output in final_res.outputs:
+            choice_data = ChatCompletionResponseChoice(
+                index=output.index,
+                message=ChatMessage(role=role, audio=audio_obj),
+                logprobs=None,
+                finish_reason="stop",
+                stop_reason=None,
+            )
+            choices.append(choice_data)
+        return choices

vllm_omni/entrypoints/stage_utils.py

@@ -1,19 +1,19 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Optional, Tuple, Union
-
-import logging
 import json
+import logging
 import os
 import pickle
 from multiprocessing import shared_memory as _shm
+from typing import Any
 
 import cloudpickle
+from omegaconf import OmegaConf
 
 logger = logging.getLogger(__name__)
 
 
-def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> None:
+def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
     """Configure per-stage CUDA visibility and current device.
 
     Behavior
@@ -26,8 +26,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
     - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string.
     """
     try:
-        selected_physical: Optional[int] = None
-        logical_idx: Optional[int] = None
+        selected_physical: int | None = None
+        logical_idx: int | None = None
 
         if isinstance(devices, str) and "," in devices:
             os.environ["CUDA_VISIBLE_DEVICES"] = devices
@@ -37,7 +37,9 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
                     selected_physical = int(toks[0])
                     logger.debug(
                         "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s",
-                        stage_id, devices, selected_physical,
+                        stage_id,
+                        devices,
+                        selected_physical,
                     )
                 except Exception as e:
                     logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
@@ -51,43 +53,70 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
                     if 0 <= logical_idx < len(mapping):
                         selected_physical = mapping[logical_idx]
                 except Exception as e:
-                    logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e)
+                    logger.debug(
+                        "[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s",
+                        stage_id,
+                        e,
+                    )
                     selected_physical = None
             if selected_physical is None:
                 selected_physical = int(logical_idx)
             os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
             logger.debug(
                 "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES to single device",
-                stage_id, logical_idx + 1, selected_physical,
+                stage_id,
+                logical_idx + 1,
+                selected_physical,
             )
         elif devices in (None, "cpu"):
-            logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
+            logger.debug(
+                "[Stage-%s] Using default device visibility (devices=%s)",
+                stage_id,
+                devices,
+            )
         else:
             selected_physical = int(str(devices))
             os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
-            logger.debug("[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical)
+            logger.debug(
+                "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)",
+                stage_id,
+                selected_physical,
+            )
 
         try:
             import torch  # noqa: WPS433
+
             if torch.cuda.is_available():
                 try:
                     torch.cuda.set_device(0)
                 except Exception as e:
-                    logger.debug("[Stage-%s] torch.cuda.set_device(0) failed: %s", stage_id, e, exc_info=True)
+                    logger.debug(
+                        "[Stage-%s] torch.cuda.set_device(0) failed: %s",
+                        stage_id,
+                        e,
+                        exc_info=True,
+                    )
                 num = torch.cuda.device_count()
                 info = []
                 for i in range(num):
                     total = torch.cuda.get_device_properties(i).total_memory
                     free, _ = torch.cuda.mem_get_info(i)
-                    info.append({
-                        "idx": i,
-                        "name": torch.cuda.get_device_name(i),
-                        "total": int(total),
-                        "free": int(free),
-                    })
+                    info.append(
+                        {
+                            "idx": i,
+                            "name": torch.cuda.get_device_name(i),
+                            "total": int(total),
+                            "free": int(free),
+                        }
+                    )
                 logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info)
         except Exception as e:
-            logger.debug("[Stage-%s] Failed to query CUDA devices: %s", stage_id, e, exc_info=True)
+            logger.debug(
+                "[Stage-%s] Failed to query CUDA devices: %s",
+                stage_id,
+                e,
+                exc_info=True,
+            )
     except Exception as e:
         logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)
 
@@ -97,7 +126,7 @@ def serialize_obj(obj: Any) -> bytes:
     return cloudpickle.dumps(obj)
 
 
-def shm_write_bytes(payload: bytes) -> Dict[str, Any]:
+def shm_write_bytes(payload: bytes) -> dict[str, Any]:
     """Write bytes into SharedMemory and return meta dict {name,size}.
 
     Caller should close the segment; the receiver should unlink.
@@ -110,11 +139,11 @@ def shm_write_bytes(payload: bytes) -> Dict[str, Any]:
     try:
         shm.close()
     except Exception as e:
-        logger.debug("Ensure parent dir failed for %s: %s", path, e)
+        logger.debug("Failed to close shared memory: %s", e)
     return meta
 
 
-def shm_read_bytes(meta: Dict[str, Any]) -> bytes:
+def shm_read_bytes(meta: dict[str, Any]) -> bytes:
     """Read bytes from SharedMemory by meta {name,size} and cleanup."""
     shm = _shm.SharedMemory(name=meta["name"])  # type: ignore[index]
     mv = memoryview(shm.buf)
@@ -141,7 +170,7 @@ def _ensure_parent_dir(path: str) -> None:
         pass
 
 
-def append_jsonl(path: str, record: Dict[str, Any]) -> None:
+def append_jsonl(path: str, record: dict[str, Any]) -> None:
     """Append a JSON record as one line to a JSONL file (best-effort).
 
     This is safe to call from multiple processes when each process writes
@@ -158,7 +187,7 @@ def append_jsonl(path: str, record: Dict[str, Any]) -> None:
         logger.exception("Failed to append JSONL to %s", path)
 
 
-def maybe_dump_to_shm(obj: Any, threshold: int) -> Tuple[bool, Any]:
+def maybe_dump_to_shm(obj: Any, threshold: int) -> tuple[bool, Any]:
     """Dump object to SHM if serialized size exceeds threshold.
 
     Returns (True, meta) when dumped; otherwise (False, original_obj).
@@ -169,7 +198,7 @@ def maybe_dump_to_shm(obj: Any, threshold: int) -> Tuple[bool, Any]:
     return False, obj
 
 
-def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) -> Any:
+def maybe_load_from_ipc(container: dict[str, Any], obj_key: str, shm_key: str) -> Any:
     """Load object from container that may carry SHM or inline object.
 
     Deprecated: prefer `maybe_load_from_ipc_with_metrics` to also obtain
@@ -180,7 +209,9 @@ def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) -
     return container[obj_key]
 
 
-def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, shm_key: str) -> tuple[Any, Dict[str, float]]:
+def maybe_load_from_ipc_with_metrics(
+    container: dict[str, Any], obj_key: str, shm_key: str
+) -> tuple[Any, dict[str, float]]:
     """Load object and return (object, metrics) with RX bytes and decode time.
 
     Metrics keys:
@@ -188,6 +219,7 @@ def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, sh
       - rx_decode_time_ms: float
     """
     import time as _time  # local import to avoid overhead at module import
+
     t0 = _time.time()
     if shm_key in container:
         meta = container[shm_key]  # type: ignore[index]
@@ -211,13 +243,13 @@ def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, sh
     }
 
 
-def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> Dict[str, Any]:
+def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> dict[str, Any]:
     """Return a dict payload for IPC: inline (obj_key) or SHM (shm_key).
 
     When serialized size exceeds threshold, returns {shm_key: {name,size}};
     otherwise returns {obj_key: obj}.
     """
-    payload: Dict[str, Any] = {}
+    payload: dict[str, Any] = {}
     use_shm, data = maybe_dump_to_shm(obj, threshold)
     if use_shm:
         payload[shm_key] = data
@@ -227,7 +259,7 @@ def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> Dict
 
 
 # Convert OmegaConf/objects to plain dicts
-def _to_dict(x: Any) -> Dict[str, Any]:
+def _to_dict(x: Any) -> dict[str, Any]:
     try:
         if isinstance(x, dict):
             return dict(x)

vllm_omni/entrypoints/utils.py

@@ -1,10 +1,14 @@
+from __future__ import annotations
+
+import logging
 import os
 from pathlib import Path
 
 from omegaconf import OmegaConf
-
 from vllm.transformers_utils.config import get_config
 
+logger = logging.getLogger(__name__)
+
 # Get the project root directory (2 levels up from this file)
 PROJECT_ROOT = Path(__file__).parent.parent.parent
 

vllm_omni/inputs/data.py

@@ -1,7 +1,12 @@
-from typing import Any, NotRequired, Optional
+from typing import Any, Optional
+
+try:
+    from typing import NotRequired
+except ImportError:
+    # Python < 3.11: use typing_extensions
+    from typing_extensions import NotRequired
 
 import torch
-
 from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt
 
 
@@ -39,8 +44,7 @@ def token_inputs_omni(
     prompt_embeds: Optional[torch.Tensor] = None,
     additional_information: Optional[dict[str, Any]] = None,
 ) -> OmniTokenInputs:
-    """Construct [`TokenInputs`][vllm.inputs.data.TokenInputs] from optional
-    values."""
+    """Construct token inputs with optional embeddings and metadata."""
     inputs = OmniTokenInputs(type="token", prompt_token_ids=prompt_token_ids)
 
     if prompt is not None:

vllm_omni/inputs/parse.py

@@ -17,15 +17,9 @@ def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPromp
         # Priority tokens: When both tokens and embeds exist, keep both and
         # follow the tokens path
         if "prompt_token_ids" in prompt:
-            return ParsedTokensPrompt(
-                type="tokens", content=prompt
-            )  # type: ignore[typeddict-item]
+            return ParsedTokensPrompt(type="tokens", content=prompt)  # type: ignore[typeddict-item]
         elif "prompt_embeds" in prompt:
-            return ParsedEmbedsPrompt(
-                type="embeds", content=prompt
-            )  # type: ignore[typeddict-item]
+            return ParsedEmbedsPrompt(type="embeds", content=prompt)  # type: ignore[typeddict-item]
         elif "prompt" in prompt:
             return ParsedTextPrompt(type="text", content=prompt)
-    raise TypeError(
-        "inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt"
-    )
+    raise TypeError("inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt")

vllm_omni/inputs/preprocess.py

@@ -1,13 +1,16 @@
 from typing import Any, Optional, Union
 
 from typing_extensions import assert_never
-
 from vllm.inputs.data import SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt
 from vllm.inputs.preprocess import InputPreprocessor
+from vllm.logger import init_logger
 from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict
+
 from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni
 from vllm_omni.inputs.parse import parse_singleton_prompt_omni
 
+logger = init_logger(__name__)
+
 
 class OmniInputPreprocessor(InputPreprocessor):
     def _process_tokens(
@@ -17,9 +20,7 @@ class OmniInputPreprocessor(InputPreprocessor):
         *,
         mm_uuids: Optional[MultiModalUUIDDict] = None,
     ) -> Union[OmniTokenInputs, MultiModalInputs]:
-        prompt_token_ids = self._truncate_inputs(
-            parsed_content["prompt_token_ids"], tokenization_kwargs
-        )
+        prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs)
         prompt_embeds = parsed_content.get("prompt_embeds")
         additional_information = parsed_content.get("additional_information")
 
@@ -62,7 +63,7 @@ class OmniInputPreprocessor(InputPreprocessor):
 
         Returns:
 
-        * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance
+        * Input container compatible with vLLM's singleton prompt handling.
         """
         parsed = parse_singleton_prompt_omni(prompt)
 

vllm_omni/model_executor/layers/mrope.py

@@ -4,7 +4,6 @@ from typing import Optional, Union
 import numpy as np
 import torch
 from transformers import PretrainedConfig
-
 from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding
 
 
@@ -249,9 +248,7 @@ class MRotaryEmbedding(RotaryEmbedding):
                     input_token_type.append("text")
 
             input_type_group: list[tuple[str, int, int]] = []
-            for key, group_iter in itertools.groupby(
-                enumerate(input_token_type), lambda x: x[1]
-            ):
+            for key, group_iter in itertools.groupby(enumerate(input_token_type), lambda x: x[1]):
                 group_list = list(group_iter)
                 start_index = group_list[0][0]
                 end_index = group_list[-1][0] + 1
@@ -260,9 +257,7 @@ class MRotaryEmbedding(RotaryEmbedding):
             video_frame_num = 1
             mm_data_idx = 0
             for modality_type, start_idx, end_idx in input_type_group:
-                st_idx = (
-                    llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-                )
+                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
                 if modality_type == "image":
                     t, h, w = (
                         image_grid_thw[mm_data_idx][0],
@@ -275,27 +270,10 @@ class MRotaryEmbedding(RotaryEmbedding):
                         w // spatial_merge_size,
                     )
 
-                    t_index = (
-                        torch.arange(llm_grid_t)
-                        .view(-1, 1)
-                        .expand(-1, llm_grid_h * llm_grid_w)
-                        .flatten()
-                    )
-                    h_index = (
-                        torch.arange(llm_grid_h)
-                        .view(1, -1, 1)
-                        .expand(llm_grid_t, -1, llm_grid_w)
-                        .flatten()
-                    )
-                    w_index = (
-                        torch.arange(llm_grid_w)
-                        .view(1, 1, -1)
-                        .expand(llm_grid_t, llm_grid_h, -1)
-                        .flatten()
-                    )
-                    llm_pos_ids_list.append(
-                        torch.stack([t_index, h_index, w_index]) + st_idx
-                    )
+                    t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
+                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+                    llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx)
                     mm_data_idx += 1
 
                 elif modality_type == "video":
@@ -311,36 +289,17 @@ class MRotaryEmbedding(RotaryEmbedding):
                     )
 
                     for t_idx in range(llm_grid_t):
-                        t_index = (
-                            torch.tensor(t_idx)
-                            .view(-1, 1)
-                            .expand(-1, llm_grid_h * llm_grid_w)
-                            .flatten()
-                        )
-                        h_index = (
-                            torch.arange(llm_grid_h)
-                            .view(1, -1, 1)
-                            .expand(1, -1, llm_grid_w)
-                            .flatten()
-                        )
-                        w_index = (
-                            torch.arange(llm_grid_w)
-                            .view(1, 1, -1)
-                            .expand(1, llm_grid_h, -1)
-                            .flatten()
-                        )
-                        llm_pos_ids_list.append(
-                            torch.stack([t_index, h_index, w_index]) + st_idx
-                        )
+                        t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
+                        h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten()
+                        w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten()
+                        llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx)
 
                     mm_data_idx += 1
                     video_frame_num += 1
 
                 else:
                     text_len = end_idx - start_idx
-                    llm_pos_ids_list.append(
-                        torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
-                    )
+                    llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
                     video_frame_num = 1
 
         else:
@@ -372,9 +331,7 @@ class MRotaryEmbedding(RotaryEmbedding):
         tokens_per_second = getattr(hf_config.vision_config, "tokens_per_second", 1.0)
 
         input_tokens_tensor = torch.tensor(input_tokens)
-        vision_start_indices = torch.argwhere(
-            input_tokens_tensor == vision_start_token_id
-        ).squeeze(1)
+        vision_start_indices = torch.argwhere(input_tokens_tensor == vision_start_token_id).squeeze(1)
         vision_tokens = input_tokens_tensor[vision_start_indices + 1]
         image_nums = (vision_tokens == image_token_id).sum()
         video_nums = (vision_tokens == video_token_id).sum()
@@ -424,15 +381,11 @@ class MRotaryEmbedding(RotaryEmbedding):
             text_len = ed - st
 
             st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-            llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
-            )
+            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
 
             t_index = (
                 (
-                    torch.arange(llm_grid_t)
-                    .view(-1, 1)
-                    .expand(-1, llm_grid_h * llm_grid_w)
+                    torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w)
                     * video_second_per_grid_t
                     * tokens_per_second
                 )
@@ -440,29 +393,15 @@ class MRotaryEmbedding(RotaryEmbedding):
                 .flatten()
             )
 
-            h_index = (
-                torch.arange(llm_grid_h)
-                .view(1, -1, 1)
-                .expand(llm_grid_t, -1, llm_grid_w)
-                .flatten()
-            )
-            w_index = (
-                torch.arange(llm_grid_w)
-                .view(1, 1, -1)
-                .expand(llm_grid_t, llm_grid_h, -1)
-                .flatten()
-            )
-            llm_pos_ids_list.append(
-                torch.stack([t_index, h_index, w_index]) + text_len + st_idx
-            )
+            h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+            w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+            llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
             st = ed + llm_grid_t * llm_grid_h * llm_grid_w
 
         if st < len(input_tokens):
             st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
             text_len = len(input_tokens) - st
-            llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
-            )
+            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
 
         llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
         mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
@@ -512,9 +451,7 @@ class MRotaryEmbedding(RotaryEmbedding):
         vision_end_token_id = thinker_config.vision_end_token_id
         seconds_per_chunk = thinker_config.seconds_per_chunk
         spatial_merge_size = thinker_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(
-            thinker_config.vision_config, "tokens_per_second", 25
-        )
+        tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
 
         if isinstance(image_grid_thw, list):
             image_grid_thw = torch.tensor(image_grid_thw)
@@ -534,21 +471,13 @@ class MRotaryEmbedding(RotaryEmbedding):
         idx = 0
         while idx < len(src_item):
             new_src_item_len = len(new_src_item)
-            start_idx = (
-                llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
-            )
+            start_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
             if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]:
                 if use_audio_in_video and idx > 0:
-                    if (
-                        src_item[idx] == vision_end_token_id
-                        and src_item[idx - 1] == audio_end_token_id
-                    ):
+                    if src_item[idx] == vision_end_token_id and src_item[idx - 1] == audio_end_token_id:
                         # processing the <|audio_eos|> before <|vision_eos|>
                         start_idx -= 1
-                    elif (
-                        src_item[idx] == audio_start_token_id
-                        and src_item[idx - 1] == vision_start_token_id
-                    ):
+                    elif src_item[idx] == audio_start_token_id and src_item[idx - 1] == vision_start_token_id:
                         # processing the <|audio_bos|> after <|vision_eos|>
                         start_idx -= 1
                 new_src_item.append(src_item[idx])
@@ -571,58 +500,40 @@ class MRotaryEmbedding(RotaryEmbedding):
                     start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
                 )
                 llm_pos_ids_list.append(llm_pos_ids)
-                vision_seqlen = image_grid_thw[image_idx].prod() // (
-                    spatial_merge_size**2
-                )
+                vision_seqlen = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
                 new_src_item.extend([image_token_id] * vision_seqlen)
                 image_idx += 1
             elif src_item[idx] == video_token_id and not use_audio_in_video:
                 grid_t = video_grid_thw[video_idx][0]
                 grid_hs = video_grid_thw[:, 1]
                 grid_ws = video_grid_thw[:, 2]
-                t_index = (
-                    torch.arange(grid_t)
-                    * second_per_grid_ts[video_idx]
-                    * tokens_per_second
-                ).long()
+                t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
                 llm_pos_ids = cls._get_llm_pos_ids_for_vision(
                     start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
                 )
                 llm_pos_ids_list.append(llm_pos_ids)
-                vision_seqlen = video_grid_thw[video_idx].prod() // (
-                    spatial_merge_size**2
-                )
+                vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
                 new_src_item.extend([video_token_id] * vision_seqlen)
                 video_idx += 1
             else:
                 # read audio from video
                 assert audio_seqlens is not None
                 audio_seqlen = audio_seqlens[audio_idx]
-                vision_seqlen = video_grid_thw[video_idx].prod() // (
-                    spatial_merge_size**2
-                )
+                vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
                 grid_t = video_grid_thw[video_idx][0]
                 grid_h = video_grid_thw[video_idx][1]
                 grid_w = video_grid_thw[video_idx][2]
                 grid_hs = video_grid_thw[:, 1]
                 grid_ws = video_grid_thw[:, 2]
                 t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
-                t_index = (
-                    torch.arange(grid_t)
-                    * second_per_grid_ts[video_idx]
-                    * tokens_per_second
-                ).long()
-                t_index_split_chunk = cls._split_list_into_ranges(
-                    t_index, t_ntoken_per_chunk
-                )
+                t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
+                t_index_split_chunk = cls._split_list_into_ranges(t_index, t_ntoken_per_chunk)
                 place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
                 pure_audio_len = place_num - 2
                 added_audio_len = 0
                 audio_llm_pos_ids_list: list[torch.Tensor] = []
                 for t_chunk in t_index_split_chunk:
-                    vision_ntoken_per_chunk = (
-                        len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
-                    )
+                    vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
                     new_src_item.extend([video_token_id] * vision_ntoken_per_chunk)
                     vision_llm_pos_ids_list = cls._get_llm_pos_ids_for_vision(
                         start_idx,
@@ -633,38 +544,23 @@ class MRotaryEmbedding(RotaryEmbedding):
                         grid_ws,
                     ).split(1, dim=1)
                     llm_pos_ids_list.extend(vision_llm_pos_ids_list)
-                    new_src_item.extend(
-                        min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
-                        * [audio_token_id]
-                    )
+                    new_src_item.extend(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) * [audio_token_id])
                     audio_start_idx = (
-                        start_idx
-                        if len(audio_llm_pos_ids_list) == 0
-                        else audio_llm_pos_ids_list[-1][0].item() + 1
+                        start_idx if len(audio_llm_pos_ids_list) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
                     )
                     if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0:
                         audio_llm_pos_ids_list = (
-                            torch.arange(
-                                min(
-                                    t_ntoken_per_chunk, pure_audio_len - added_audio_len
-                                )
-                            ).expand(3, -1)
+                            torch.arange(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)).expand(3, -1)
                             + audio_start_idx
                         ).split(1, dim=1)
                     else:
                         audio_llm_pos_ids_list = []
-                    added_audio_len += min(
-                        t_ntoken_per_chunk, pure_audio_len - added_audio_len
-                    )
+                    added_audio_len += min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
                     llm_pos_ids_list.extend(audio_llm_pos_ids_list)
                 if added_audio_len < pure_audio_len:
-                    new_src_item.extend(
-                        (pure_audio_len - added_audio_len) * [audio_token_id]
-                    )
+                    new_src_item.extend((pure_audio_len - added_audio_len) * [audio_token_id])
                     audio_llm_pos_ids_list = (
-                        torch.arange(pure_audio_len - added_audio_len).expand(3, -1)
-                        + llm_pos_ids_list[-1].max()
-                        + 1
+                        torch.arange(pure_audio_len - added_audio_len).expand(3, -1) + llm_pos_ids_list[-1].max() + 1
                     ).split(1, dim=1)
                     llm_pos_ids_list.extend(audio_llm_pos_ids_list)
                 audio_idx += 1
@@ -673,9 +569,7 @@ class MRotaryEmbedding(RotaryEmbedding):
             idx += len(new_src_item) - new_src_item_len
 
         llm_positions = torch.cat(llm_pos_ids_list, dim=1)
-        mrope_position_delta = (
-            torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item)
-        )
+        mrope_position_delta = torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item)
         llm_positions = llm_positions[:, context_len:seq_len]
 
         return llm_positions, mrope_position_delta
@@ -692,25 +586,10 @@ class MRotaryEmbedding(RotaryEmbedding):
         llm_pos_ids_list = []
         llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
         llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
-        h_index = (
-            torch.arange(llm_grid_h)
-            .view(1, -1, 1)
-            .expand(len(t_index), -1, llm_grid_w)
-            .flatten()
-        )
-        w_index = (
-            torch.arange(llm_grid_w)
-            .view(1, 1, -1)
-            .expand(len(t_index), llm_grid_h, -1)
-            .flatten()
-        )
+        h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten()
+        w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten()
         t_index_tensor = (
-            torch.Tensor(t_index)
-            .to(llm_grid_h.device)
-            .view(-1, 1)
-            .expand(-1, llm_grid_h * llm_grid_w)
-            .long()
-            .flatten()
+            torch.Tensor(t_index).to(llm_grid_h.device).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten()
         )
         _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
         llm_pos_ids_list.append(_llm_pos_ids + start_idx)
@@ -731,14 +610,7 @@ class MRotaryEmbedding(RotaryEmbedding):
         context_len: int,
         seq_len: int,
     ) -> list[list[int]]:
-        return [
-            list(
-                range(
-                    context_len + mrope_position_delta, seq_len + mrope_position_delta
-                )
-            )
-            for _ in range(3)
-        ]
+        return [list(range(context_len + mrope_position_delta, seq_len + mrope_position_delta)) for _ in range(3)]
 
     @staticmethod
     def get_next_input_positions_tensor(
@@ -748,7 +620,6 @@ class MRotaryEmbedding(RotaryEmbedding):
         context_len: int,
         num_new_tokens: int,
     ):
-
         values = np.arange(
             mrope_position_delta + context_len,
             mrope_position_delta + context_len + num_new_tokens,
@@ -779,25 +650,19 @@ class MRotaryEmbedding(RotaryEmbedding):
         audio_end_token_id = thinker_config.audio_end_token_id
         seconds_per_chunk = thinker_config.seconds_per_chunk
         spatial_merge_size = thinker_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(
-            thinker_config.vision_config, "tokens_per_second", 25
-        )
+        tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
 
         grid_t = video_grid_thw[0]
         grid_h = video_grid_thw[1]
         grid_w = video_grid_thw[2]
         t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
-        t_index = (
-            torch.arange(grid_t) * video_second_per_grid_t * tokens_per_second
-        ).long()
+        t_index = (torch.arange(grid_t) * video_second_per_grid_t * tokens_per_second).long()
         t_index_split_chunk = cls._split_list_into_ranges(t_index, t_ntoken_per_chunk)
 
         updates = [audio_start_token_id]
         added_audio_len = 0
         for t_chunk in t_index_split_chunk:
-            vision_ntoken_per_chunk = (
-                len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
-            )
+            vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
             updates.extend([video_token_id] * vision_ntoken_per_chunk)
 
             audio_chunk_size = min(t_ntoken_per_chunk, audio_len - added_audio_len)

vllm_omni/model_executor/model_loader/weight_utils.py

@@ -4,7 +4,6 @@ from typing import Optional, Union
 
 import huggingface_hub
 from huggingface_hub import snapshot_download
-
 from vllm.logger import init_logger
 from vllm.model_executor.model_loader.weight_utils import DisabledTqdm, get_lock
 

vllm_omni/model_executor/models/qwen2_5_omni.py

@@ -1,7 +1,8 @@
 import glob
 import os
+from collections.abc import Iterable
 from functools import cached_property
-from typing import Dict, Iterable, NamedTuple, Optional, Set, Tuple, Union
+from typing import NamedTuple, Optional, Union
 
 import numpy as np
 import torch
@@ -11,7 +12,6 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
     Qwen2_5OmniTalkerConfig,
     Qwen2_5OmniThinkerConfig,
 )
-
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
@@ -29,9 +29,8 @@ from vllm.sequence import IntermediateTensors
 from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
-from vllm_omni.model_executor.model_loader.weight_utils import (
-    download_weights_from_hf_specific,
-)
+
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
@@ -57,7 +56,6 @@ logger = init_logger(__name__)
 class Qwen2_5OmniForConditionalGeneration(
     nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
 ):
-
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.have_multimodal_outputs = True
@@ -117,17 +115,15 @@ class Qwen2_5OmniForConditionalGeneration(
                     architectures=["Qwen2_5OmniToken2WavModel"],
                 )
             # voice resources (loaded on demand)
-            self._token2wav_conds: Dict[str, torch.Tensor] = {}
-            self._token2wav_ref_mels: Dict[str, torch.Tensor] = {}
+            self._token2wav_conds: dict[str, torch.Tensor] = {}
+            self._token2wav_ref_mels: dict[str, torch.Tensor] = {}
             self.model = self.token2wav
         else:
             raise ValueError("Invalid model stage")
 
         # Set up intermediate tensors
         self.make_empty_intermediate_tensors = (
-            (self.thinker.make_empty_intermediate_tensors)
-            if self.model_stage == "thinker"
-            else lambda: None
+            (self.thinker.make_empty_intermediate_tensors) if self.model_stage == "thinker" else lambda: None
         )
 
     # -------------------- Device utilities --------------------
@@ -176,11 +172,7 @@ class Qwen2_5OmniForConditionalGeneration(
         multimodal_embeddings=None,
     ) -> torch.Tensor:
         if self.model_stage == "code2wav":
-            return (
-                torch.zeros_like(input_ids)
-                .reshape(-1, 1)
-                .repeat(1, self.vllm_config.model_config.get_hidden_size())
-            )
+            return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size())
         return self.model.get_input_embeddings(input_ids, multimodal_embeddings)
 
     def get_multimodal_embeddings(self, **kwargs):
@@ -229,9 +221,7 @@ class Qwen2_5OmniForConditionalGeneration(
             # if input_ids is None, set it to a zero tensor, in the length of the
             # same as the embedding seq length
             if input_ids is None:
-                input_ids = torch.zeros(
-                    inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev
-                ).unsqueeze(
+                input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev).unsqueeze(
                     0
                 )  # (1, 0)
                 added_batch_dim = True
@@ -259,11 +249,7 @@ class Qwen2_5OmniForConditionalGeneration(
 
             # Text-only path
             return OmniOutput(
-                text_hidden_states=(
-                    text_hidden_states.squeeze(0)
-                    if added_batch_dim
-                    else text_hidden_states
-                ),
+                text_hidden_states=(text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states),
                 multimodal_outputs=None,
             )
 
@@ -273,13 +259,17 @@ class Qwen2_5OmniForConditionalGeneration(
             # Rules:
             # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
             # - Decode segments consist of a single non-special token.
-            # - If additional_information is provided (can be a list split by request or a concatenated tensor plus a list of shapes),
+            # - If additional_information is provided
+            # (can be a list split by request or a concatenated tensor plus a list of shapes),
             #   then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments;
-            # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them; otherwise, keep the original embedding.
+            # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them;
+            # otherwise, keep the original embedding.
 
             if input_ids is None and additional_information is None:
                 input_ids = torch.zeros(
-                    inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device
+                    inputs_embeds.shape[0],
+                    dtype=torch.long,
+                    device=inputs_embeds.device,
                 )
                 additional_information = {}
                 self.thinker_reply_part = torch.zeros_like(inputs_embeds)
@@ -293,13 +283,17 @@ class Qwen2_5OmniForConditionalGeneration(
 
             # ------- Request-scoped additional information (no cross-request concat) -------
             request_ids: Optional[list[str]] = kwargs.get("request_ids")  # ordered
-            request_token_spans: Optional[list[tuple[int,int]]] = kwargs.get("request_token_spans")
+            request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get("request_token_spans")
             addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id")
             runtime_addi = kwargs.get("runtime_additional_information")
 
             # Normalize runtime_addi into a mapping by request_id for convenience
             runtime_addi_by_req: dict[str, dict] = {}
-            if isinstance(request_ids, list) and isinstance(runtime_addi, list) and len(runtime_addi) == len(request_ids):
+            if (
+                isinstance(request_ids, list)
+                and isinstance(runtime_addi, list)
+                and len(runtime_addi) == len(request_ids)
+            ):
                 for i, rid in enumerate(request_ids):
                     if isinstance(rid, str) and isinstance(runtime_addi[i], dict):
                         runtime_addi_by_req[rid] = runtime_addi[i]
@@ -332,9 +326,19 @@ class Qwen2_5OmniForConditionalGeneration(
                     otoks = info.get("thinker_output_token_ids")  # list[int]
 
                     if not isinstance(pe, torch.Tensor):
-                        pe = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
+                        pe = torch.zeros(
+                            0,
+                            self.talker.config.hidden_size,
+                            dtype=inputs_embeds.dtype,
+                            device=self._module_device(self.model),
+                        )
                     if not isinstance(tr, torch.Tensor):
-                        tr = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
+                        tr = torch.zeros(
+                            0,
+                            self.talker.config.hidden_size,
+                            dtype=inputs_embeds.dtype,
+                            device=self._module_device(self.model),
+                        )
                     if not isinstance(ptoks, (list, torch.Tensor)):
                         ptoks = []
                     if not isinstance(otoks, (list, torch.Tensor)):
@@ -354,7 +358,9 @@ class Qwen2_5OmniForConditionalGeneration(
 
                     # Prepare per-request reply queue for subsequent decode: drop first row
                     if tr.ndim == 2 and tr.shape[0] > 0:
-                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = tr[1:].detach().to("cpu").contiguous()
+                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = (
+                            tr[1:].detach().to("cpu").contiguous()
+                        )
 
             # ------- Decode: span_len == 1 -------
             if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list):
@@ -378,7 +384,11 @@ class Qwen2_5OmniForConditionalGeneration(
                         dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None
                         if isinstance(dv, torch.Tensor) and dv.numel() > 0:
                             step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
-                        elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0:
+                        elif (
+                            hasattr(self, "thinker_reply_part")
+                            and isinstance(self.thinker_reply_part, torch.Tensor)
+                            and self.thinker_reply_part.numel() > 0
+                        ):
                             # C) fallback shared pool
                             step_vec = self.thinker_reply_part[0:1]
                             self.thinker_reply_part = self.thinker_reply_part[1:]
@@ -403,7 +413,8 @@ class Qwen2_5OmniForConditionalGeneration(
                 multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id}
 
             if sampling_metadata is not None:
-                # the padding token id is set to text model's pad token id, which do not match with the talker model's word embedding size
+                # the padding token id is set to text model's pad token id,
+                # which do not match with the talker model's word embedding size
                 sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448
 
             return OmniOutput(
@@ -425,12 +436,8 @@ class Qwen2_5OmniForConditionalGeneration(
             code = code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code
             code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code
 
-            audio_tensor = self.generate_audio(
-                code, voice_type
-            )
-            return OmniOutput(
-                text_hidden_states=None, multimodal_outputs={"audio": audio_tensor}
-            )
+            audio_tensor = self.generate_audio(code, voice_type)
+            return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audio_tensor})
 
         return OmniOutput(
             text_hidden_states=torch.cat(
@@ -441,9 +448,7 @@ class Qwen2_5OmniForConditionalGeneration(
                     ).to(self._module_device(self.model)),
                     self.talker.thinker_to_talker_proj(
                         self.talker.get_input_embeddings(
-                            torch.tensor(
-                                [TALKER_CODEC_BOS_TOKEN_ID, TALKER_CODEC_EOS_TOKEN_ID]
-                            )
+                            torch.tensor([TALKER_CODEC_BOS_TOKEN_ID, TALKER_CODEC_EOS_TOKEN_ID])
                             .to(torch.bfloat16)
                             .to(self._module_device(self.model))
                         )
@@ -570,7 +575,6 @@ class Qwen2_5OmniForConditionalGeneration(
         thinker_prompt_embeds,
         prompt_token_ids,
     ):
-
         talker_hf_config = self.talker_config
         if hasattr(talker_hf_config, "talker_config"):
             talker_hf_config = talker_hf_config.talker_config
@@ -593,9 +597,7 @@ class Qwen2_5OmniForConditionalGeneration(
         input_tokens_len = len(prompt_token_ids_processed)
         # the code below is from model runner in Qwen, may need to further discuss later
         if input_tokens_len > 2:
-            prompt_token_ids_processed = [
-                self.talker_config.tts_codec_mask_token_id
-            ] * (input_tokens_len - 2) + [
+            prompt_token_ids_processed = [self.talker_config.tts_codec_mask_token_id] * (input_tokens_len - 2) + [
                 self.talker_config.tts_codec_pad_token_id,
                 self.talker_config.tts_codec_start_token_id,
             ]
@@ -606,9 +608,7 @@ class Qwen2_5OmniForConditionalGeneration(
             ][-input_tokens_len:]
         if isinstance(prompt_token_ids_processed, list):
             prompt_token_ids_processed = (
-                torch.Tensor(prompt_token_ids_processed)
-                .to(torch.int64)
-                .to(self._module_device(self.talker))
+                torch.Tensor(prompt_token_ids_processed).to(torch.int64).to(self._module_device(self.talker))
             )
         return prompt_token_ids_processed, prompt_embeds
 
@@ -617,14 +617,12 @@ class Qwen2_5OmniForConditionalGeneration(
         output_prompt_embeds,
         output_token_ids,
     ):
-        processed_output_token_embeds = (
-            output_prompt_embeds + self.talker.get_input_embeddings(output_token_ids)
+        processed_output_token_embeds = output_prompt_embeds + self.talker.get_input_embeddings(
+            output_token_ids
         )  # for decode
         return output_token_ids, processed_output_token_embeds
 
-    def compute_logits(
-        self, hidden_states: Union[torch.Tensor, OmniOutput]
-    ) -> Optional[torch.Tensor]:
+    def compute_logits(self, hidden_states: Union[torch.Tensor, OmniOutput]) -> Optional[torch.Tensor]:
         # Handle OmniOutput type
         if isinstance(hidden_states, OmniOutput):
             hidden_states = hidden_states.text_hidden_states
@@ -640,7 +638,7 @@ class Qwen2_5OmniForConditionalGeneration(
         # Use thinker model for sampling
         return self.model.sample(logits, sampling_metadata)
 
-    def generate_speech(self, text_tokens: torch.Tensor, voice_type: str = "default"):
+    def generate_speech(self, text_tokens: torch.Tensor, voice_type: str = "default") -> torch.Tensor:
         """
         Generate speech from text tokens using the talker and token2wav models.
         This method is kept for backward compatibility and direct speech generation.
@@ -653,9 +651,7 @@ class Qwen2_5OmniForConditionalGeneration(
             Audio tensor
         """
         # Generate codec tokens using talker model
-        talker_output = self.talker(
-            input_ids=None, positions=None, inputs_embeds=text_tokens
-        )
+        talker_output = self.talker(input_ids=None, positions=None, inputs_embeds=text_tokens)
 
         # Convert talker output to codec tokens
         codec_tokens = self._convert_to_codec_tokens(talker_output)
@@ -682,9 +678,7 @@ class Qwen2_5OmniForConditionalGeneration(
             # Suppress only codec_bos, consistent with HF generate's
             # suppress_tokens behavior
             bos_id = None
-            if hasattr(self, "talker_config") and hasattr(
-                self.talker_config, "tts_codec_start_token_id"
-            ):
+            if hasattr(self, "talker_config") and hasattr(self.talker_config, "tts_codec_start_token_id"):
                 bos_id = int(getattr(self.talker_config, "tts_codec_start_token_id"))
             if bos_id is not None:
                 logits[..., bos_id] = -1e9
@@ -703,12 +697,8 @@ class Qwen2_5OmniForConditionalGeneration(
         conds = getattr(self.token2wav_config, "conds", None)
         ref_mels = getattr(self.token2wav_config, "ref_mels", None)
         if isinstance(conds, dict) and isinstance(ref_mels, dict):
-            self._token2wav_conds = {
-                k: torch.as_tensor(v, device=device) for k, v in conds.items()
-            }
-            self._token2wav_ref_mels = {
-                k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()
-            }
+            self._token2wav_conds = {k: torch.as_tensor(v, device=device) for k, v in conds.items()}
+            self._token2wav_ref_mels = {k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()}
         # legacy: load from directory if provided
         model_path = hf_model_folder
         if isinstance(model_path, str) and os.path.isdir(model_path):
@@ -720,24 +710,14 @@ class Qwen2_5OmniForConditionalGeneration(
                     self._token2wav_ref_mels[key] = value["ref_mel"].to(device)
             else:
                 # legacy npy inputs
-                for f in sorted(
-                    glob.glob(os.path.join(model_path, "inputs", "*spk_emb.npy"))
-                ):
+                for f in sorted(glob.glob(os.path.join(model_path, "inputs", "*spk_emb.npy"))):
                     key = os.path.basename(f).split("_")[0].lower()
-                    self._token2wav_conds[key] = torch.as_tensor(
-                        np.load(f), device=device
-                    )
-                for f in sorted(
-                    glob.glob(os.path.join(model_path, "inputs", "*ref_mel.npy"))
-                ):
+                    self._token2wav_conds[key] = torch.as_tensor(np.load(f), device=device)
+                for f in sorted(glob.glob(os.path.join(model_path, "inputs", "*ref_mel.npy"))):
                     key = os.path.basename(f).split("_")[0].lower()
-                    self._token2wav_ref_mels[key] = torch.as_tensor(
-                        np.load(f), device=device
-                    )
+                    self._token2wav_ref_mels[key] = torch.as_tensor(np.load(f), device=device)
 
-    def _codec_to_audio(
-        self, codec_tokens: torch.Tensor, voice_type: str = "default"
-    ) -> Optional[torch.Tensor]:
+    def _codec_to_audio(self, codec_tokens: torch.Tensor, voice_type: str = "default") -> Optional[torch.Tensor]:
         if self.token2wav is None:
             self._init_token2wav_model()
         if self.token2wav is None:
@@ -771,9 +751,7 @@ class Qwen2_5OmniForConditionalGeneration(
             if codec.ndim == 1:
                 codec = codec.unsqueeze(0)
         else:
-            codec = torch.as_tensor(
-                codec_tokens, dtype=torch.long, device=token2wav_dev
-            ).unsqueeze(0)
+            codec = torch.as_tensor(codec_tokens, dtype=torch.long, device=token2wav_dev).unsqueeze(0)
 
         # Streaming with chunked process and boundary alignment
         # (rely on token2wav.process_chunk)
@@ -788,9 +766,7 @@ class Qwen2_5OmniForConditionalGeneration(
         steps = 10
 
         # Prepare initial noise for the whole sequence
-        y_all = torch.randn(
-            (1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev
-        )
+        y_all = torch.randn((1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev)
 
         logger.info(
             "Currently, we do not use the chunked process, we only use the "
@@ -802,9 +778,7 @@ class Qwen2_5OmniForConditionalGeneration(
         for i in range(codec.shape[1]):
             chunk_code_length = i * 2 - 24
             finished = i == (codec.shape[1] - 1)
-            if (
-                chunk_code_length > 0 and chunk_code_length % chunk_size == 0
-            ) or finished:
+            if (chunk_code_length > 0 and chunk_code_length % chunk_size == 0) or finished:
                 chunk_ends.append(i)
 
         # Number of chunks in mel domain
@@ -833,7 +807,7 @@ class Qwen2_5OmniForConditionalGeneration(
         waveform = np.concatenate(wav_chunks)
         return torch.as_tensor(waveform, device=token2wav_dev)
 
-    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights for all components of the omni model."""
         loaded_weights = set()
         thinker_weights = []
@@ -862,20 +836,14 @@ class Qwen2_5OmniForConditionalGeneration(
         if talker_weights and self.talker is not None:
             # Map talker weights to appropriate components
             if self.thinker is None:
-                thinker_embedding_weights = [
-                    w
-                    for n, w in thinker_weights
-                    if n == "thinker.model.embed_tokens.weight"
-                ]
+                thinker_embedding_weights = [w for n, w in thinker_weights if n == "thinker.model.embed_tokens.weight"]
                 if thinker_embedding_weights:
                     self.thinker_embedding = nn.Embedding(
                         thinker_embedding_weights[0].shape[0],
                         thinker_embedding_weights[0].shape[1],
                     )
                     self.thinker_embedding.weight = nn.Parameter(
-                        thinker_embedding_weights[0].to(
-                            self._module_device(self.talker)
-                        )
+                        thinker_embedding_weights[0].to(self._module_device(self.talker))
                     )
             talker_loaded = self.talker.load_weights(talker_weights)
             talker_loaded = add_prefix_to_loaded_weights(talker_loaded, "talker")
@@ -896,9 +864,7 @@ class Qwen2_5OmniForConditionalGeneration(
                     allow_patterns=["*.pt"],
                 )
             self._init_token2wav_model(hf_model_folder)
-            t2w_loaded = self.token2wav.load_weights(
-                token2wav_weights, os.path.join(hf_model_folder, "spk_dict.pt")
-            )
+            t2w_loaded = self.token2wav.load_weights(token2wav_weights, os.path.join(hf_model_folder, "spk_dict.pt"))
             t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, "token2wav")
             loaded_weights.update(t2w_loaded)
 

vllm_omni/model_executor/models/qwen2_5_omni_talker.py

@@ -1,24 +1,17 @@
+from collections.abc import Iterable
 from functools import cached_property
-from typing import Iterable, Optional, Set, Tuple, Union
+from typing import Optional, Union
 
 import torch
 import torch.nn as nn
-from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
-    Qwen2_5OmniTalkerConfig,
-)
-from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
-    Qwen2_5OmniAudioEncoder,
-)
+from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniTalkerConfig
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
 
 # from vllm.attention import AttentionMetadata  # unused import
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ColumnParallelLinear
-from vllm.model_executor.models.interfaces import (
-    MultiModalEmbeddings,
-    SupportsMultiModal,
-    SupportsPP,
-)
+from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
     Qwen2_5OmniThinkerDummyInputsBuilder,
@@ -72,9 +65,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
 
         if hasattr(config, "talker_config"):
             self.config = config.talker_config
-            vllm_config.model_config.hf_text_config = (
-                vllm_config.model_config.hf_config.talker_config
-            )
+            vllm_config.model_config.hf_text_config = vllm_config.model_config.hf_config.talker_config
         else:
             self.config = config
 
@@ -92,9 +83,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
             hf_config=getattr(self.config, "text_config", self.config),
             architectures=["Qwen2ForCausalLM_old"],
         )
-        self.make_empty_intermediate_tensors = (
-            self.language_model.make_empty_intermediate_tensors
-        )
+        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
 
     def init_multi_modal(self, thinker_config):
         self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
@@ -122,7 +111,6 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.get_input_embeddings(input_ids)
         if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
-
             # TODO (ywang96): support overlapping modalitiy embeddings so that
             # `use_audio_in_video` will work on V1.
             inputs_embeds = merge_multimodal_embeddings(
@@ -145,9 +133,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         inputs_embeds: Optional[torch.Tensor] = None,
         **kwargs: object,
     ) -> Union[torch.Tensor, IntermediateTensors]:
-        assert (
-            input_ids is not None or inputs_embeds is not None
-        ), "input_ids or inputs_embeds must be provided"
+        assert input_ids is not None or inputs_embeds is not None, "input_ids or inputs_embeds must be provided"
         # forward_context: ForwardContext = get_forward_context()  # unused variable
 
         if intermediate_tensors is not None:
@@ -166,8 +152,16 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         )
         return hidden_states
 
+    def bad_word_processor(self, logits: torch.Tensor) -> torch.Tensor:
+        if hasattr(self.config, "tts_codec_start_token_id"):
+            bos_id = int(getattr(self.config, "tts_codec_start_token_id"))
+            logits[..., bos_id] = -1e9
+        return logits
+
     def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:
-        return self.language_model.compute_logits(hidden_states)
+        logits = self.language_model.compute_logits(hidden_states)
+        logits = self.bad_word_processor(logits)
+        return logits
 
     def sample(
         self,
@@ -176,7 +170,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
     ) -> Optional[SamplerOutput]:
         return self.language_model.sample(logits, sampling_metadata)
 
-    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(
             self,
             skip_prefixes=["thinker.", "token2wav."],
@@ -212,31 +206,15 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         # Preserve the order of modalities if there are multiple of them
         # from the order of kwargs.
         for input_key in kwargs:
-            if (
-                input_key in ("pixel_values", "image_embeds")
-                and "image" not in mm_input_by_modality
-            ):
-                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
-                    **kwargs
-                )
-            if (
-                input_key in ("pixel_values_videos", "video_embeds")
-                and "video" not in mm_input_by_modality
-            ):
-                mm_input_by_modality["video"] = self._parse_and_validate_video_input(
-                    **kwargs
-                )
-            if (
-                input_key in ("input_audio_features")
-                and "audio" not in mm_input_by_modality
-            ):
-                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
-                    **kwargs
-                )
+            if input_key in ("pixel_values", "image_embeds") and "image" not in mm_input_by_modality:
+                mm_input_by_modality["image"] = self._parse_and_validate_image_input(**kwargs)
+            if input_key in ("pixel_values_videos", "video_embeds") and "video" not in mm_input_by_modality:
+                mm_input_by_modality["video"] = self._parse_and_validate_video_input(**kwargs)
+            if input_key in ("input_audio_features") and "audio" not in mm_input_by_modality:
+                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs)
         return mm_input_by_modality
 
     def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
-
         mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []

vllm_omni/model_executor/models/qwen2_5_omni_thinker.py

@@ -6,25 +6,13 @@ from typing import Any, Optional, Union
 import torch
 import torch.nn as nn
 from transformers.feature_extraction_utils import BatchFeature
-from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
-    Qwen2_5OmniConfig,
-    Qwen2_5OmniThinkerConfig,
-)
-from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
-    Qwen2_5OmniAudioEncoder,
-)
-from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import (
-    Qwen2_5OmniProcessor,
-)
+from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
+from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import Qwen2_5OmniProcessor
 from transformers.models.whisper import WhisperFeatureExtractor
-
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import (
-    MultiModalEmbeddings,
-    SupportsMultiModal,
-    SupportsPP,
-)
+from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VisionTransformer,
     Qwen2_5_VLImageEmbeddingInputs,
@@ -64,15 +52,11 @@ from vllm.multimodal.parse import (
     MultiModalDataItems,
     MultiModalDataParser,
 )
-from vllm.multimodal.processing import (
-    BaseMultiModalProcessor,
-    PlaceholderFeaturesInfo,
-    PromptReplacement,
-    PromptUpdate,
-)
+from vllm.multimodal.processing import BaseMultiModalProcessor, PlaceholderFeaturesInfo, PromptReplacement, PromptUpdate
 from vllm.multimodal.profiling import BaseDummyInputsBuilder
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.tokenizer import decode_tokens, encode_tokens
+
 from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
 
 try:
@@ -93,17 +77,13 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
     video_grid_sizes = video_grid_thw.prod(-1)
 
     return dict(
-        input_audio_features=MultiModalFieldConfig.flat_from_sizes(
-            "audio", audio_feature_lengths, dim=1
-        ),
+        input_audio_features=MultiModalFieldConfig.flat_from_sizes("audio", audio_feature_lengths, dim=1),
         feature_attention_mask=MultiModalFieldConfig.batched("audio"),
         audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
         pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
         image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
         image_grid_thw=MultiModalFieldConfig.batched("image"),
-        pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
-            "video", video_grid_sizes
-        ),
+        pixel_values_videos=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes),
         video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes),
         video_grid_thw=MultiModalFieldConfig.batched("video"),
         second_per_grid_ts=MultiModalFieldConfig.batched("video"),
@@ -111,7 +91,6 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
 
 
 class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
-
     def _parse_audio_data(
         self,
         data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],
@@ -127,10 +106,7 @@ class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
         return super()._parse_audio_data(data)
 
 
-class Qwen2_5OmniThinkerProcessingInfo(
-    Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo
-):
-
+class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo):
     def get_hf_config(self):
         return self.ctx.get_hf_config(Qwen2_5OmniConfig).thinker_config
 
@@ -151,10 +127,7 @@ class Qwen2_5OmniThinkerProcessingInfo(
         return {"audio": None, "image": None, "video": None}
 
 
-class Qwen2_5OmniThinkerDummyInputsBuilder(
-    BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]
-):
-
+class Qwen2_5OmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]):
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
         num_audios = mm_counts.get("audio", 0)
         num_images = mm_counts.get("image", 0)
@@ -166,11 +139,7 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
         image_token: str = hf_processor.image_token
         video_token: str = hf_processor.video_token
 
-        return (
-            audio_token * num_audios
-            + image_token * num_images
-            + video_token * num_videos
-        )
+        return audio_token * num_audios + image_token * num_images + video_token * num_videos
 
     def get_dummy_mm_data(
         self,
@@ -191,17 +160,11 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
             * feature_extractor.sampling_rate
         )
         target_width, target_height = self.info.get_image_size_with_most_features()
-        target_num_frames = self.info.get_num_frames_with_most_features(
-            seq_len, mm_counts
-        )
+        target_num_frames = self.info.get_num_frames_with_most_features(seq_len, mm_counts)
 
         mm_data = {
-            "audio": self._get_dummy_audios(
-                length=target_audio_length, num_audios=num_audios
-            ),
-            "image": self._get_dummy_images(
-                width=target_width, height=target_height, num_images=num_images
-            ),
+            "audio": self._get_dummy_audios(length=target_audio_length, num_audios=num_audios),
+            "image": self._get_dummy_images(width=target_width, height=target_height, num_images=num_images),
             "video": self._get_dummy_videos(
                 width=target_width,
                 height=target_height,
@@ -213,15 +176,10 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
         return mm_data
 
 
-class Qwen2_5OmniThinkerMultiModalProcessor(
-    BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]
-):
-
+class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]):
     def _get_data_parser(self) -> MultiModalDataParser:
         feature_extractor = self.info.get_feature_extractor()
-        return Qwen2_5OmniThinkerMultiModalDataParser(
-            target_sr=feature_extractor.sampling_rate
-        )
+        return Qwen2_5OmniThinkerMultiModalDataParser(target_sr=feature_extractor.sampling_rate)
 
     def _call_hf_processor(
         self,
@@ -252,14 +210,9 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         feature_attention_mask = hf_inputs.get("feature_attention_mask", None)
         if "input_audio_features" not in hf_inputs and input_features is not None:
             if feature_attention_mask is not None:
-                input_features = input_features.permute(0, 2, 1)[
-                    feature_attention_mask.bool()
-                ].permute(1, 0)
+                input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
             hf_inputs["input_audio_features"] = input_features
-        if (
-            "audio_feature_lengths" not in hf_inputs
-            and feature_attention_mask is not None
-        ):
+        if "audio_feature_lengths" not in hf_inputs and feature_attention_mask is not None:
             hf_inputs["audio_feature_lengths"] = feature_attention_mask.sum(-1)
         return hf_inputs
 
@@ -299,9 +252,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 prompt_ids,
                 mm_item_counts,
             )
-            self._validate_mm_placeholders(
-                mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
-            )
+            self._validate_mm_placeholders(mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video)
 
             tokenizer = self.info.get_tokenizer()
             prompt = decode_tokens(tokenizer, prompt_ids)
@@ -315,9 +266,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 mm_prompt_updates,
                 mm_item_counts,
             )
-            self._validate_mm_placeholders(
-                mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
-            )
+            self._validate_mm_placeholders(mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video)
 
         tokenizer = self.info.get_tokenizer()
         prompt = decode_tokens(tokenizer, prompt_ids)
@@ -350,15 +299,11 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         if audio_feature_lengths is None and feature_attention_mask is None:
             audio_output_lengths = []
         elif audio_feature_lengths is not None:
-            _, audio_output_lens = _get_feat_extract_output_lengths(
-                audio_feature_lengths
-            )
+            _, audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths)
             audio_output_lengths = audio_output_lens.tolist()
         elif feature_attention_mask is not None:
             assert isinstance(feature_attention_mask, torch.Tensor)
-            _, audio_output_lens = _get_feat_extract_output_lengths(
-                feature_attention_mask.sum(-1)
-            )
+            _, audio_output_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1))
             audio_output_lengths = audio_output_lens.tolist()
 
         # number of audios read from video.
@@ -372,8 +317,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 audios = mm_items.get_items("audio", AudioProcessorItems)
                 audio = audios.get(item_idx)
                 raise ValueError(
-                    f"The audio {audio} (len={len(audio)}) is too short "
-                    "to be represented inside the model"
+                    f"The audio {audio} (len={len(audio)}) is too short to be represented inside the model"
                 )
 
             return [audio_token_id] * num_features
@@ -392,9 +336,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         def get_replacement_qwen2_use_audio_in_video(item_idx: int):
             nonlocal audio_in_video_item_idx
 
-            audio_num_features = audio_output_lengths[
-                audio_in_video_item_idx + item_idx
-            ]
+            audio_num_features = audio_output_lengths[audio_in_video_item_idx + item_idx]
             video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx]
 
             audio_in_video_item_idx += 1
@@ -509,39 +451,27 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
 
 
 class Qwen2_5OmniConditionalGenerationMixin:
-
-    def _validate_and_reshape_mm_tensor(
-        self, mm_input: object, name: str, dim: int = 0
-    ) -> torch.Tensor:
+    def _validate_and_reshape_mm_tensor(self, mm_input: object, name: str, dim: int = 0) -> torch.Tensor:
         if not isinstance(mm_input, (torch.Tensor, list)):
-            raise ValueError(
-                f"Incorrect type of {name}. " f"Got type: {type(mm_input)}"
-            )
+            raise ValueError(f"Incorrect type of {name}. Got type: {type(mm_input)}")
         if isinstance(mm_input, torch.Tensor):
             return torch.concat(list(mm_input), dim=dim)
         else:
             return torch.concat(mm_input, dim=dim)
 
-    def _parse_and_validate_audio_input(
-        self, **kwargs: object
-    ) -> Optional[Qwen2AudioInputs]:
+    def _parse_and_validate_audio_input(self, **kwargs: object) -> Optional[Qwen2AudioInputs]:
         input_audio_features = kwargs.pop("input_audio_features", None)
         audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
         feature_attention_mask = kwargs.pop("feature_attention_mask", None)
         if input_audio_features is None:
             return None
-        input_audio_features = self._validate_and_reshape_mm_tensor(
-            input_audio_features, "input_audio_features", dim=1
-        )
+        input_audio_features = self._validate_and_reshape_mm_tensor(input_audio_features, "input_audio_features", dim=1)
         if feature_attention_mask is not None:
             feature_attention_mask = self._validate_and_reshape_mm_tensor(
                 feature_attention_mask, "feature_attention_mask"
             )
         if not isinstance(input_audio_features, (torch.Tensor, list)):
-            raise ValueError(
-                "Incorrect type of audio input features. "
-                f"Got type: {type(input_audio_features)}"
-            )
+            raise ValueError(f"Incorrect type of audio input features. Got type: {type(input_audio_features)}")
         return Qwen2AudioInputs(
             input_features=input_audio_features,
             audio_feature_lengths=audio_feature_lengths,
@@ -560,18 +490,11 @@ class Qwen2_5OmniConditionalGenerationMixin:
             return None
 
         if pixel_values is not None:
-            pixel_values = self._validate_and_reshape_mm_tensor(
-                pixel_values, "image pixel values"
-            )
-            image_grid_thw = self._validate_and_reshape_mm_tensor(
-                image_grid_thw, "image grid_thw"
-            )
+            pixel_values = self._validate_and_reshape_mm_tensor(pixel_values, "image pixel values")
+            image_grid_thw = self._validate_and_reshape_mm_tensor(image_grid_thw, "image grid_thw")
 
             if not isinstance(pixel_values, (torch.Tensor, list)):
-                raise ValueError(
-                    "Incorrect type of image pixel values. "
-                    f"Got type: {type(pixel_values)}"
-                )
+                raise ValueError(f"Incorrect type of image pixel values. Got type: {type(pixel_values)}")
 
             return Qwen2_5_VLImagePixelInputs(
                 type="pixel_values",
@@ -580,18 +503,11 @@ class Qwen2_5OmniConditionalGenerationMixin:
             )
 
         if image_embeds is not None:
-            image_embeds = self._validate_and_reshape_mm_tensor(
-                image_embeds, "image embeds"
-            )
-            image_grid_thw = self._validate_and_reshape_mm_tensor(
-                image_grid_thw, "image grid_thw"
-            )
+            image_embeds = self._validate_and_reshape_mm_tensor(image_embeds, "image embeds")
+            image_grid_thw = self._validate_and_reshape_mm_tensor(image_grid_thw, "image grid_thw")
 
             if not isinstance(image_embeds, torch.Tensor):
-                raise ValueError(
-                    "Incorrect type of image embeddings. "
-                    f"Got type: {type(image_embeds)}"
-                )
+                raise ValueError(f"Incorrect type of image embeddings. Got type: {type(image_embeds)}")
             return Qwen2_5_VLImageEmbeddingInputs(
                 type="image_embeds",
                 image_embeds=image_embeds,
@@ -610,12 +526,8 @@ class Qwen2_5OmniConditionalGenerationMixin:
             return None
 
         if pixel_values_videos is not None:
-            pixel_values_videos = self._validate_and_reshape_mm_tensor(
-                pixel_values_videos, "video pixel values"
-            )
-            video_grid_thw = self._validate_and_reshape_mm_tensor(
-                video_grid_thw, "video grid_thw"
-            )
+            pixel_values_videos = self._validate_and_reshape_mm_tensor(pixel_values_videos, "video pixel values")
+            video_grid_thw = self._validate_and_reshape_mm_tensor(video_grid_thw, "video grid_thw")
 
             return Qwen2_5_VLVideoPixelInputs(
                 type="pixel_values_videos",
@@ -624,18 +536,11 @@ class Qwen2_5OmniConditionalGenerationMixin:
             )
 
         if video_embeds is not None:
-            video_embeds = self._validate_and_reshape_mm_tensor(
-                video_embeds, "video embeds"
-            )
-            video_grid_thw = self._validate_and_reshape_mm_tensor(
-                video_grid_thw, "video grid_thw"
-            )
+            video_embeds = self._validate_and_reshape_mm_tensor(video_embeds, "video embeds")
+            video_grid_thw = self._validate_and_reshape_mm_tensor(video_grid_thw, "video grid_thw")
 
             if not isinstance(video_embeds, torch.Tensor):
-                raise ValueError(
-                    "Incorrect type of video embeddings. "
-                    f"Got type: {type(video_embeds)}"
-                )
+                raise ValueError(f"Incorrect type of video embeddings. Got type: {type(video_embeds)}")
             return Qwen2_5_VLVideoEmbeddingInputs(
                 type="video_embeds",
                 video_embeds=video_embeds,
@@ -648,24 +553,20 @@ class Qwen2_5OmniConditionalGenerationMixin:
         audio_hashes: list[str] = None,
         cached_audio_features: torch.Tensor = None,
     ) -> torch.Tensor:
-
         input_features = audio_input["input_features"]
         audio_feature_lengths = audio_input["audio_feature_lengths"]
         if input_features.ndim == 3:
             assert input_features.shape[0] == 1
             input_features = input_features.squeeze(0)
         if audio_feature_lengths.ndim == 2:
-            assert (
-                audio_feature_lengths.shape[0] == 1
-                or audio_feature_lengths.shape[1] == 1
-            )
+            assert audio_feature_lengths.shape[0] == 1 or audio_feature_lengths.shape[1] == 1
             if audio_feature_lengths.shape[0] == 1:
                 audio_feature_lengths = audio_feature_lengths.squeeze(0)
             else:
                 audio_feature_lengths = audio_feature_lengths.squeeze(1)
 
-        audio_feat_lengths, audio_output_lengths = (
-            self.audio_tower._get_feat_extract_output_lengths(audio_feature_lengths)
+        audio_feat_lengths, audio_output_lengths = self.audio_tower._get_feat_extract_output_lengths(
+            audio_feature_lengths
         )
 
         audio_outputs = self.audio_tower(
@@ -676,9 +577,7 @@ class Qwen2_5OmniConditionalGenerationMixin:
         audio_features = audio_outputs.last_hidden_state
         return audio_features.split(audio_output_lengths.tolist())
 
-    def _process_image_input(
-        self, image_input: Qwen2_5_VLImageInputs
-    ) -> tuple[torch.Tensor, ...]:
+    def _process_image_input(self, image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]:
         if image_input["type"] == "image_embeds":
             return image_input["image_embeds"].type(self.visual.dtype)
 
@@ -758,8 +657,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         else:
             logger.warning(
                 "flash_attn is not available, the model may not yield the "
-                "exactly same result as the transformers implementation "
-                "in the audio tower part."
+                "exactly same result as the transformers implementation in the audio tower part."
             )
 
         self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
@@ -777,9 +675,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             architectures=["Qwen2ForCausalLM"],
         )
 
-        self.make_empty_intermediate_tensors = (
-            self.language_model.make_empty_intermediate_tensors
-        )
+        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
 
     def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
         mm_input_by_modality = {}
@@ -787,34 +683,18 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         # Preserve the order of modalities if there are multiple of them
         # from the order of kwargs.
         for input_key in kwargs:
-            if (
-                input_key in ("pixel_values", "image_embeds")
-                and "image" not in mm_input_by_modality
-            ):
-                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
-                    **kwargs
-                )
-            if (
-                input_key in ("pixel_values_videos", "video_embeds")
-                and "video" not in mm_input_by_modality
-            ):
-                mm_input_by_modality["video"] = self._parse_and_validate_video_input(
-                    **kwargs
-                )
-            if (
-                input_key in ("input_audio_features")
-                and "audio" not in mm_input_by_modality
-            ):
-                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
-                    **kwargs
-                )
+            if input_key in ("pixel_values", "image_embeds") and "image" not in mm_input_by_modality:
+                mm_input_by_modality["image"] = self._parse_and_validate_image_input(**kwargs)
+            if input_key in ("pixel_values_videos", "video_embeds") and "video" not in mm_input_by_modality:
+                mm_input_by_modality["video"] = self._parse_and_validate_video_input(**kwargs)
+            if input_key in ("input_audio_features") and "audio" not in mm_input_by_modality:
+                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs)
         return mm_input_by_modality
 
     def get_language_model(self) -> torch.nn.Module:
         return self.language_model
 
     def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
-
         mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []
@@ -845,7 +725,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.get_input_embeddings(input_ids)
         if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
-
             # TODO (ywang96): support overlapping modalitiy embeddings so that
             # `use_audio_in_video` will work on V1.
             inputs_embeds = merge_multimodal_embeddings(
@@ -897,9 +776,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
                 placeholder_token_id = self.config.image_token_index
             if modality == "video":
                 placeholder_token_id = self.config.video_token_index
-            inputs_embeds = merge_multimodal_embeddings(
-                input_ids, inputs_embeds, embeddings, placeholder_token_id
-            )
+            inputs_embeds = merge_multimodal_embeddings(input_ids, inputs_embeds, embeddings, placeholder_token_id)
         return inputs_embeds
 
     def forward(
@@ -921,10 +798,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             text_inputs_embeds = self.get_input_embeddings(
                 input_ids,
                 (
-                    [
-                        (torch.zeros_like(embeddings), modality)
-                        for embeddings, modality in multimodal_embeddings
-                    ]
+                    [(torch.zeros_like(embeddings), modality) for embeddings, modality in multimodal_embeddings]
                     if multimodal_embeddings is not None
                     else None
                 ),

vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py

@@ -3,7 +3,8 @@
 ############################
 
 import math
-from typing import Iterable, List, Optional, Set, Tuple, Union
+from collections.abc import Iterable
+from typing import Optional, Union
 
 import numpy as np
 import torch
@@ -16,29 +17,25 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
     Qwen2_5OmniDiTConfig,
     Qwen2_5OmniToken2WavConfig,
 )
-from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
-    Qwen2_5OmniPreTrainedModel,
-)
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniPreTrainedModel
 
 # Bring in HF base classes, configs and utilities used below
 from transformers.utils.logging import get_logger as _hf_get_logger
-
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
+from vllm.model_executor.model_loader.utils import set_default_torch_dtype
 from vllm.model_executor.models.interfaces import SupportsPP
-from vllm.model_executor.models.utils import (
-    AutoWeightsLoader as _Vllm_AutoWeightsLoader,
-)
+from vllm.model_executor.models.utils import AutoWeightsLoader as _Vllm_AutoWeightsLoader
 from vllm.model_executor.models.utils import WeightsMapper as _Vllm_WeightsMapper
-from vllm.model_executor.models.utils import (
-    init_vllm_registered_model as _Vllm_init_vllm_registered_model,
-)
-from vllm.model_executor.models.utils import maybe_prefix as _Vllm_maybe_prefix
+from vllm.model_executor.models.utils import init_vllm_registered_model as _vllm_init_vllm_registered_model
+from vllm.model_executor.models.utils import maybe_prefix as _vllm_maybe_prefix
 from vllm.sequence import IntermediateTensors
 from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.model_executor.ode_solver import get_ode_solver_class
+
 
 # Provide a no-op auto_docstring decorator to satisfy annotations if missing
 def auto_docstring(func=None, **_kwargs):
@@ -218,9 +215,7 @@ class AttentiveStatisticsPooling(nn.Module):
 
     def _compute_statistics(self, x, m, dim=2):
         mean = (m * x).sum(dim)
-        std = torch.sqrt(
-            (m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps)
-        )
+        std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps))
         return mean, std
 
     def forward(self, hidden_states):
@@ -282,9 +277,7 @@ class SqueezeExcitationRes2NetBlock(nn.Module):
             kernel_size=1,
             dilation=1,
         )
-        self.res2net_block = Res2NetBlock(
-            out_channels, out_channels, res2net_scale, kernel_size, dilation
-        )
+        self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation)
         self.tdnn2 = TimeDelayNetBlock(
             out_channels,
             out_channels,
@@ -312,13 +305,10 @@ class ECAPA_TimeDelayNet(torch.nn.Module):
 
     def __init__(self, config: Qwen2_5OmniDiTConfig):
         super().__init__()
-        if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(
-            config.enc_channels
-        ) != len(config.enc_dilations):
-            raise ValueError(
-                "enc_channels, enc_kernel_sizes and enc_dilations should have "
-                "same length"
-            )
+        if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(config.enc_channels) != len(
+            config.enc_dilations
+        ):
+            raise ValueError("enc_channels, enc_kernel_sizes and enc_dilations should have same length")
         self.channels = config.enc_channels
         self.blocks = nn.ModuleList()
 
@@ -412,26 +402,14 @@ class DiTInputEmbedding(nn.Module):
     ):
         if apply_cfg:
             hidden_states = torch.cat([hidden_states, hidden_states], dim=0)
-            speaker_embedding = torch.cat(
-                [speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0
-            )
-            condition_vector = torch.cat(
-                [condition_vector, torch.zeros_like(condition_vector)], dim=0
-            )
+            speaker_embedding = torch.cat([speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0)
+            condition_vector = torch.cat([condition_vector, torch.zeros_like(condition_vector)], dim=0)
             code_embed = torch.cat([code_embed, code_embed_uncond], dim=0)
         elif drop_audio_cond:  # cfg for cond audio
             condition_vector = torch.zeros_like(condition_vector)
             speaker_embedding = torch.zeros_like(speaker_embedding)
-        condition_vector = (
-            self.spk_encoder(condition_vector)
-            .unsqueeze(1)
-            .repeat(1, hidden_states.size(1), 1)
-        )
-        hidden_states = self.proj(
-            torch.cat(
-                (hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1
-            )
-        )
+        condition_vector = self.spk_encoder(condition_vector).unsqueeze(1).repeat(1, hidden_states.size(1), 1)
+        hidden_states = self.proj(torch.cat((hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1))
 
         return hidden_states
 
@@ -465,13 +443,9 @@ class Qwen2_5_OmniAdaLayerNormZero(nn.Module):
 
     def forward(self, hidden_states, emb=None):
         emb = self.linear(self.silu(emb))
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(
-            emb, 6, dim=1
-        )
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
 
-        hidden_states = (
-            self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
-        )
+        hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
         return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp
 
 
@@ -490,9 +464,7 @@ class Qwen2_5_OmniAdaLayerNormZero_Final(nn.Module):
         emb = self.linear(self.silu(emb))
         scale, shift = torch.chunk(emb, 2, dim=1)
 
-        hidden_states = (
-            self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
-        )
+        hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
         return hidden_states
 
 
@@ -572,9 +544,7 @@ class DiTAttention(nn.Module):
         self.to_k = nn.Linear(config.hidden_size, self.inner_dim)
         self.to_v = nn.Linear(config.hidden_size, self.inner_dim)
 
-        self.to_out = nn.ModuleList(
-            [nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)]
-        )
+        self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)])
 
     def forward(
         self,
@@ -600,11 +570,10 @@ class DiTAttention(nn.Module):
         # Due to training process, only first head is applied with RoPE,
         # will be fixed at next release
         cos, sin = position_embeddings
-        query[:, :1], key[:, :1] = apply_rotary_pos_emb(
-            query[:, :1], key[:, :1], cos, sin
-        )
+        query[:, :1], key[:, :1] = apply_rotary_pos_emb(query[:, :1], key[:, :1], cos, sin)
 
-        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        attn_impl = self.config._attn_implementation or "sdpa"
+        attention_interface = ALL_ATTENTION_FUNCTIONS[attn_impl]
         attention_weights, _ = attention_interface(
             self,
             query,
@@ -616,9 +585,7 @@ class DiTAttention(nn.Module):
 
         # mask. e.g. inference got a batch with different target durations,
         # mask out the padding
-        attention_weights = attention_weights.reshape(
-            batch_size, -1, self.heads * head_dim
-        )
+        attention_weights = attention_weights.reshape(batch_size, -1, self.heads * head_dim)
         attention_weights = attention_weights.to(query.dtype)
 
         # linear proj
@@ -648,9 +615,7 @@ class DiTTimestepEmbedding(nn.Module):
     def __init__(self, dim, freq_embed_dim=256):
         super().__init__()
         self.time_embed = SinusPositionEmbedding(freq_embed_dim)
-        self.time_mlp = nn.ModuleList(
-            [nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)]
-        )
+        self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)])
 
     def forward(self, timestep):  # noqa: F821
         time_hidden = self.time_embed(timestep)
@@ -661,29 +626,21 @@ class DiTTimestepEmbedding(nn.Module):
 
 
 class DiTDecoderLayer(nn.Module):
-    def __init__(
-        self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0
-    ):
+    def __init__(self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0):
         super().__init__()
         self.attn_norm = Qwen2_5_OmniAdaLayerNormZero(config.hidden_size)
 
         self.attn = DiTAttention(config)
         self.look_ahead_block = look_ahead_block
         self.look_backward_block = look_backward_block
-        self.ff_norm = nn.LayerNorm(
-            config.hidden_size, elementwise_affine=False, eps=1e-6
-        )
-        self.ff = DiTMLP(
-            dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout
-        )
+        self.ff_norm = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=1e-6)
+        self.ff = DiTMLP(dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout)
 
     def forward(
         self, hidden_states, timestep, position_embeddings=None, block_diff=None
     ):  # x: noised input, t: time embedding
         # pre-norm & modulation for attention input
-        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(
-            hidden_states, emb=timestep
-        )
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(hidden_states, emb=timestep)
 
         # attention
         attn_output = self.attn(
@@ -696,9 +653,7 @@ class DiTDecoderLayer(nn.Module):
         # process attention output for input x
         hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_output
 
-        norm = (
-            self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
-        )
+        norm = self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
         ff_output = self.ff(norm)
         hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
 
@@ -741,14 +696,14 @@ class SnakeBeta(nn.Module):
         beta = self.beta.unsqueeze(0).unsqueeze(-1)
         alpha = torch.exp(alpha)
         beta = torch.exp(beta)
-        hidden_states = hidden_states + (
-            1.0 / (beta + self.no_div_by_zero)
-        ) * torch.pow(torch.sin(hidden_states * alpha), 2)
+        hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
+            torch.sin(hidden_states * alpha), 2
+        )
 
         return hidden_states
 
 
-def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
+def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
     """Generates a 1D Kaiser-windowed sinc filter.
 
     Args:
@@ -773,9 +728,7 @@ def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
     else:
         beta = 0.0
 
-    kaiser_window = torch.kaiser_window(
-        kernel_size, beta=beta, periodic=False, dtype=torch.float32
-    )
+    kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
 
     # Compute time indices
     if is_even:
@@ -785,9 +738,7 @@ def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
 
     # Compute sinc filter
     if cutoff == 0:
-        return torch.zeros(
-            (1, 1, kernel_size), dtype=torch.float32
-        )  # Ensures correct shape
+        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)  # Ensures correct shape
 
     sinc_filter = torch.sinc(2 * cutoff * time_indices)
     normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
@@ -802,27 +753,19 @@ class UpSample1d(nn.Module):
     def __init__(self, ratio=2, kernel_size=None):
         super().__init__()
         self.ratio = ratio
-        self.kernel_size = (
-            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
-        )
+        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
         self.stride = ratio
         self.pad = self.kernel_size // ratio - 1
         self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
-        self.pad_right = (
-            self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
-        )
+        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
 
-        filter = kaiser_sinc_filter1d(
-            cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
-        )
+        filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
         self.register_buffer("filter", filter, persistent=False)
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
         hidden_states_dtype = hidden_states.dtype
-        hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(
-            self.filter.dtype
-        )
+        hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype)
         hidden_states = self.ratio * F.conv_transpose1d(
             hidden_states,
             self.filter.expand(channels, -1, -1),
@@ -855,9 +798,7 @@ class DownSample1d(nn.Module):
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
         hidden_states_dtype = hidden_states.dtype
-        hidden_states = F.pad(
-            hidden_states, (self.pad_left, self.pad_right), mode="replicate"
-        ).to(self.filter.dtype)
+        hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(self.filter.dtype)
         out = F.conv1d(
             hidden_states,
             self.filter.expand(channels, -1, -1),
@@ -958,15 +899,10 @@ class AMPBlock(torch.nn.Module):
             ]
         )
 
-        self.num_layers = len(self.convs1) + len(
-            self.convs2
-        )  # total number of conv layers
+        self.num_layers = len(self.convs1) + len(self.convs2)  # total number of conv layers
 
         self.activations = nn.ModuleList(
-            [
-                TorchActivation1d(activation=SnakeBeta(channels))
-                for _ in range(self.num_layers)
-            ]
+            [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
         )
 
     def _get_padding(self, kernel_size, dilation=1):
@@ -999,9 +935,7 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
         self.num_residual_blocks = len(config.resblock_kernel_sizes)
         self.num_upsample_layers = len(config.upsample_rates)
 
-        self.conv_pre = nn.Conv1d(
-            config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3
-        )
+        self.conv_pre = nn.Conv1d(config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3)
 
         # Removing extra ModuleList breaks official state dict
         ups = [
@@ -1016,9 +950,7 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
                     )
                 ]
             )
-            for layer_idx, (stride, kernel_size) in enumerate(
-                zip(config.upsample_rates, config.upsample_kernel_sizes)
-            )
+            for layer_idx, (stride, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes))
         ]
         self.ups = nn.ModuleList(ups)
 
@@ -1030,16 +962,12 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
                     dilation,
                 )
                 for layer_idx in range(self.num_upsample_layers)
-                for kernel_size, dilation in zip(
-                    config.resblock_kernel_sizes, config.resblock_dilation_sizes
-                )
+                for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes)
             ]
         )
 
         self.activation_post = TorchActivation1d(
-            activation=SnakeBeta(
-                config.upsample_initial_channel // (2**self.num_upsample_layers)
-            )
+            activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers))
         )
         self.conv_post = nn.Conv1d(
             config.upsample_initial_channel // (2**self.num_upsample_layers),
@@ -1079,9 +1007,7 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
         for layer_index in range(self.num_upsample_layers):
             hidden_representation = self.ups[layer_index][0](hidden_representation)
             residual_output = sum(
-                self.resblocks[layer_index * self.num_residual_blocks + block_index](
-                    hidden_representation
-                )
+                self.resblocks[layer_index * self.num_residual_blocks + block_index](hidden_representation)
                 for block_index in range(self.num_residual_blocks)
             )
             residual_output = residual_output / self.num_residual_blocks
@@ -1092,118 +1018,24 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
         return torch.clamp(output_waveform, min=-1.0, max=1.0).squeeze().cpu()
 
 
-class RungeKutta4ODESolver:
-    def __init__(self, function, initial_value):
-        self.function = function
-        self.initial_value = initial_value
-
-        self._one_third = 1 / 3
-        self._two_thirds = 2 / 3
-
-    def _rk4_step(
-        self,
-        function,
-        time_start,
-        time_step,
-        time_end,
-        value_start,
-        function_value_start=None,
-    ):
-        k1 = (
-            function_value_start
-            if function_value_start is not None
-            else function(time_start, value_start)
-        )
-        k2 = function(
-            time_start + time_step * self._one_third,
-            value_start + time_step * k1 * self._one_third,
-        )
-        k3 = function(
-            time_start + time_step * self._two_thirds,
-            value_start + time_step * (k2 - k1 * self._one_third),
-        )
-        k4 = function(time_end, value_start + time_step * (k1 - k2 + k3))
-        return (k1 + 3 * (k2 + k3) + k4) * time_step / 8
-
-    def _compute_step(self, function, time_start, time_step, time_end, value_start):
-        function_value_start = function(time_start, value_start)
-        return (
-            self._rk4_step(
-                function,
-                time_start,
-                time_step,
-                time_end,
-                value_start,
-                function_value_start=function_value_start,
-            ),
-            function_value_start,
-        )
-
-    def _linear_interpolation(
-        self, time_start, time_end, value_start, value_end, time_point
-    ):
-        if time_point == time_start:
-            return value_start
-        if time_point == time_end:
-            return value_end
-        weight = (time_point - time_start) / (time_end - time_start)
-        return value_start + weight * (value_end - value_start)
-
-    def integrate(self, time_points):
-        solution = torch.empty(
-            len(time_points),
-            *self.initial_value.shape,
-            dtype=self.initial_value.dtype,
-            device=self.initial_value.device,
-        )
-        solution[0] = self.initial_value
-
-        current_index = 1
-        current_value = self.initial_value
-        for time_start, time_end in zip(time_points[:-1], time_points[1:]):
-            time_step = time_end - time_start
-            delta_value, _ = self._compute_step(
-                self.function, time_start, time_step, time_end, current_value
-            )
-            next_value = current_value + delta_value
-
-            while (
-                current_index < len(time_points)
-                and time_end >= time_points[current_index]
-            ):
-                solution[current_index] = self._linear_interpolation(
-                    time_start,
-                    time_end,
-                    current_value,
-                    next_value,
-                    time_points[current_index],
-                )
-                current_index += 1
-
-            current_value = next_value
-
-        return solution
-
-
 @auto_docstring(
     custom_intro="""
     The full Qwen2.5Omni Token2WavDiT model. Which take speech tokens as
     input and predict mel spectrogram.
     """
 )
-class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
-    config: Qwen2_5OmniDiTConfig
-    _no_split_modules = ["DiTDecoderLayer"]
+class Qwen2_5OmniToken2WavDiTModel(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        model_config = vllm_config.model_config
+        config: Qwen2_5OmniDiTConfig = model_config.hf_config
 
-    def __init__(self, config: Qwen2_5OmniDiTConfig):
-        super().__init__(config)
         self.mel_dim = config.mel_dim
         self.repeats = config.repeats
         self.time_embed = DiTTimestepEmbedding(config.hidden_size)
 
-        self.text_embed = DiTCodecEmbedding(
-            config.num_embeds, config.emb_dim, config.repeats
-        )
+        self.text_embed = DiTCodecEmbedding(config.num_embeds, config.emb_dim, config.repeats)
         self.input_embed = DiTInputEmbedding(config)
 
         self.rotary_embed = Qwen2_5OmniDiTRotaryEmbedding(config.head_dim)
@@ -1223,16 +1055,16 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
                 )
             )
 
-        self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(
-            config.hidden_size
-        )  # final modulation
+        self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(config.hidden_size)  # final modulation
         self.proj_out = nn.Linear(config.hidden_size, config.mel_dim)
 
+        self.ode_solver = get_ode_solver_class(vllm_config.model_config.ode_solver_class)(
+            function=None, initial_value=None
+        )  # placeholder
+
     def _create_block_diff(self, hidden_states):
         batch, seq_len = hidden_states.shape[0], hidden_states.shape[1]
-        block_indices = (
-            torch.arange(seq_len, device=hidden_states.device) // self.block_size
-        )  # [seq_length]
+        block_indices = torch.arange(seq_len, device=hidden_states.device) // self.block_size  # [seq_length]
 
         block_i = block_indices.unsqueeze(1)  # [seq_length, 1]
         block_j = block_indices.unsqueeze(0)  # [1, seq_length]
@@ -1257,12 +1089,8 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
 
         # Compute embeddings
         time_embedding = self.time_embed(time_step)
-        text_embedding = self.text_embed(
-            quantized_code, drop_code=False if apply_cfg else drop_code
-        )
-        text_embedding_unconditioned = (
-            self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
-        )
+        text_embedding = self.text_embed(quantized_code, drop_code=False if apply_cfg else drop_code)
+        text_embedding_unconditioned = self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
 
         hidden_states = self.input_embed(
             hidden_states,
@@ -1302,17 +1130,11 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         guidance_scale=0.5,
         sway_coefficient=-1.0,
     ):
-        noise_initialization = torch.randn(
-            [1, 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype
-        )
+        noise_initialization = torch.randn([1, 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype)
         maximum_duration = quantized_code.shape[1] * self.repeats
-        initial_state = noise_initialization[:, :maximum_duration].to(
-            quantized_code.device
-        )
+        initial_state = noise_initialization[:, :maximum_duration].to(quantized_code.device)
         batch_size = reference_mel_spectrogram.shape[0]
-        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(
-            1, maximum_duration, 1
-        )
+        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, maximum_duration, 1)
 
         if batch_size != 1:
             raise ValueError("Only batch size = 1 is currently supported")
@@ -1339,10 +1161,7 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
                 apply_cfg=True,
             )
             guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
-            return (
-                guided_prediction
-                + (guided_prediction - null_prediction) * guidance_scale
-            )
+            return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
 
         initial_time = 0
         time_embedding = torch.linspace(
@@ -1354,14 +1173,10 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         )
 
         if sway_coefficient is not None:
-            time_embedding += sway_coefficient * (
-                torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding
-            )
+            time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
 
-        ode_solver = RungeKutta4ODESolver(
-            function=ode_function, initial_value=initial_state
-        )
-        solution_trajectory = ode_solver.integrate(time_embedding)
+        self.ode_solver.reset_ode_solver(function=ode_function, initial_value=initial_state)
+        solution_trajectory = self.ode_solver.integrate(time_embedding)
 
         generated_waveform = solution_trajectory[-1]
         generated_mel_spectrogram = generated_waveform.permute(0, 2, 1)
@@ -1370,14 +1185,14 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
     @torch.no_grad()
     def fast_block_sample(
         self,
-        conditioning_vector,
-        reference_mel_spectrogram,
-        quantized_code,
+        conditioning_vector: torch.Tensor,
+        reference_mel_spectrogram: torch.Tensor,
+        quantized_code: torch.Tensor,
         y0: torch.Tensor,
-        num_steps=10,
-        guidance_scale=0.5,
+        num_steps: int = 10,
+        guidance_scale: float = 0.5,
         sway_coefficient: Optional[float] = -1.0,
-    ):
+    ) -> torch.Tensor:
         """
         Block-wise ODE sampling starting from provided initial state y0.
 
@@ -1391,9 +1206,7 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         """
         initial_state = y0.to(quantized_code.device)
         batch_size = reference_mel_spectrogram.shape[0]
-        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(
-            1, initial_state.shape[1], 1
-        )
+        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, initial_state.shape[1], 1)
 
         if batch_size != 1:
             raise ValueError("Only batch size = 1 is currently supported")
@@ -1420,10 +1233,7 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
                 apply_cfg=True,
             )
             guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
-            return (
-                guided_prediction
-                + (guided_prediction - null_prediction) * guidance_scale
-            )
+            return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
 
         initial_time = 0
         time_embedding = torch.linspace(
@@ -1435,14 +1245,10 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         )
 
         if sway_coefficient is not None:
-            time_embedding += sway_coefficient * (
-                torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding
-            )
+            time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
 
-        ode_solver = RungeKutta4ODESolver(
-            function=ode_function, initial_value=initial_state
-        )
-        solution_trajectory = ode_solver.integrate(time_embedding)
+        self.ode_solver.reset_ode_solver(function=ode_function, initial_value=initial_state)
+        solution_trajectory = self.ode_solver.integrate(time_embedding)
 
         generated_waveform = solution_trajectory[-1]
         generated_mel_spectrogram = generated_waveform.permute(0, 2, 1)
@@ -1456,16 +1262,13 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
     mel spectrogram as input and predict waveform.
     """
 )
-class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
-    config: Qwen2_5OmniToken2WavConfig
-    base_model_prefix = "model"
-    _no_split_modules = [
-        "Qwen2_5OmniToken2WavDiTModel",
-        "Qwen2_5OmniToken2WavBigVGANModel",
-    ]
+class Qwen2_5OmniToken2WavModel(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix=""):
+        super().__init__()
+
+        self.model_config = vllm_config.model_config
+        config: Qwen2_5OmniToken2WavConfig = self.model_config.hf_config
 
-    def __init__(self, config: Qwen2_5OmniToken2WavConfig):
-        super().__init__(config)
         attn_impl = config._attn_implementation
         if config._attn_implementation == "flash_attention_2":
             logger.warning_once(
@@ -1477,13 +1280,13 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
             attn_impl = "sdpa"
         elif config._attn_implementation == "eager":
             logger.warning_once(
-                "Qwen2_5OmniToken2WavModel does not support eager attention "
-                "implementation, fall back to sdpa"
+                "Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa"
             )
             attn_impl = "sdpa"
-        self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel._from_config(
-            config.dit_config, attn_implementation=attn_impl
-        )
+        with set_default_torch_dtype(torch.float32):
+            self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel(
+                vllm_config=vllm_config.with_hf_config(config.dit_config),
+            )
         self.code2wav_bigvgan_model = Qwen2_5OmniToken2WavBigVGANModel._from_config(
             config.bigvgan_config, attn_implementation=attn_impl
         )
@@ -1501,9 +1304,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
 
         # codec embedding size for masking EOS out-of-range
         try:
-            self.codec_embed_size = (
-                self.code2wav_dit_model.text_embed.codec_embed.weight.size(0)
-            )
+            self.codec_embed_size = self.code2wav_dit_model.text_embed.codec_embed.weight.size(0)
         except Exception:
             self.codec_embed_size = -1
 
@@ -1518,6 +1319,10 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
             # fallback to commonly used value
             self.vocoder_hop = 240
 
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
     def forward(
         self,
         code,
@@ -1585,7 +1390,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         steps: int,
         prev_generated: torch.Tensor,
         finished: bool = False,
-    ) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
+    ) -> tuple[Optional[torch.Tensor], torch.Tensor]:
         """Streaming per small chunk: returns (mel_or_None, audio_slice)."""
         start_index = max(i * self.chunk_size - self.past_cache_size, 0)
         end_index = min(
@@ -1593,16 +1398,8 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
             codec_all.shape[1] * self.factor,
         )
 
-        y0 = (
-            y_all[:, start_index:end_index]
-            .reshape(1, -1, self.code2wav_dit_model.mel_dim)
-            .contiguous()
-        )
-        codec = (
-            codec_all[:, start_index // self.factor : end_index // self.factor]
-            .reshape(1, -1)
-            .contiguous()
-        )
+        y0 = y_all[:, start_index:end_index].reshape(1, -1, self.code2wav_dit_model.mel_dim).contiguous()
+        codec = codec_all[:, start_index // self.factor : end_index // self.factor].reshape(1, -1).contiguous()
 
         # generate mel for current window (B, mel_dim, T)
         generated = self.process_chunk_dit_batch(
@@ -1632,9 +1429,9 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         y_all: torch.Tensor,
         i: int,
         steps: int,
-        prev_generated: Union[torch.Tensor, List[torch.Tensor]],
+        prev_generated: Union[torch.Tensor, list[torch.Tensor]],
         finished: bool = False,
-    ) -> Tuple[Union[torch.Tensor, List[torch.Tensor]], torch.Tensor]:
+    ) -> tuple[Union[torch.Tensor, list[torch.Tensor]], torch.Tensor]:
         """High-level chunk API aligning to qwen2_code2wav_dit signature."""
         if not isinstance(prev_generated, torch.Tensor):
             prev_generated = prev_generated[0] if len(prev_generated) > 0 else None
@@ -1659,7 +1456,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         finished: bool,
         prev_generated: Optional[torch.Tensor],
         generated: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Align mel and audio boundaries for 50Hz-like streaming.
 
@@ -1676,21 +1473,13 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
             mel = generated[:, :, : self.chunk_size]
         elif finished:
             mel_trim = generated[:, :, self.past_cache_size :]
-            mel = torch.cat(
-                [prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2
-            )
+            mel = torch.cat([prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2)
         else:
             if start_index == 0:
-                mel_trim = generated[
-                    :, :, i * self.chunk_size : -self.future_cache_size
-                ]
+                mel_trim = generated[:, :, i * self.chunk_size : -self.future_cache_size]
             else:
-                mel_trim = generated[
-                    :, :, self.past_cache_size : -self.future_cache_size
-                ]
-            mel = torch.cat(
-                [prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2
-            )
+                mel_trim = generated[:, :, self.past_cache_size : -self.future_cache_size]
+            mel = torch.cat([prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2)
 
         audio = self.code2wav_bigvgan_model(mel)
         if i == 0:
@@ -1698,11 +1487,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         elif finished:
             audio_output = audio[self.future_size * self.vocoder_hop :]
         else:
-            audio_output = audio[
-                self.future_size
-                * self.vocoder_hop : -self.future_size
-                * self.vocoder_hop
-            ]
+            audio_output = audio[self.future_size * self.vocoder_hop : -self.future_size * self.vocoder_hop]
         return mel, audio_output
 
 
@@ -1726,9 +1511,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         self.config = vllm_config.model_config.hf_config
 
         # Initialize underlying HF Token2Wav model via registry
-        self.token2wav = _Vllm_init_vllm_registered_model(
+        self.token2wav = _vllm_init_vllm_registered_model(
             vllm_config=vllm_config,
-            prefix=_Vllm_maybe_prefix(prefix, "token2wav_model"),
+            prefix=_vllm_maybe_prefix(prefix, "token2wav_model"),
             hf_config=self.config,
             architectures=["Qwen2_5OmniToken2WavDiTModel"],
         )
@@ -1780,9 +1565,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
     ) -> Optional[SamplerOutput]:
         return None
 
-    def load_weights_without_buffers(
-        self, weights: Iterable[Tuple[str, torch.Tensor]]
-    ) -> Set[str]:
+    def load_weights_without_buffers(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = _Vllm_AutoWeightsLoader(self)
         loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
         # Log load summary
@@ -1820,9 +1603,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         return registers
 
     # remove buffers from the weights and reload them after loading weights
-    def remove_buffers_from_weights(
-        self, weights: Iterable[Tuple[str, torch.Tensor]], buffers: dict
-    ):
+    def remove_buffers_from_weights(self, weights: Iterable[tuple[str, torch.Tensor]], buffers: dict):
         weights_to_load = []
         for key, value in weights:
             if key in buffers:
@@ -1842,9 +1623,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
                 loaded_buffers.add(name)
         return loaded_buffers
 
-    def load_weights(
-        self, weights: Iterable[Tuple[str, torch.Tensor]], spk_dict_path: str
-    ) -> Set[str]:
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]], spk_dict_path: str) -> set[str]:
         buffers = self.find_all_registers()
         weights_to_load = self.remove_buffers_from_weights(weights, buffers)
         loaded = self.load_weights_without_buffers(weights_to_load)
@@ -1872,9 +1651,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         )
 
     @torch.inference_mode()
-    def process_chunk_bigvgan_batch(
-        self, mel_batch: torch.Tensor
-    ) -> Optional[torch.Tensor]:
+    def process_chunk_bigvgan_batch(self, mel_batch: torch.Tensor) -> Optional[torch.Tensor]:
         # BigVGAN is not part of this wrapper; return None for parity.
         return None
 
@@ -1889,7 +1666,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         steps: int,
         prev_generated: torch.Tensor,
         finished: bool = False,
-    ) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
+    ) -> tuple[Optional[torch.Tensor], torch.Tensor]:
         mel = self.token2wav(
             code=codec_all,
             conditioning=conditioning,
@@ -1907,9 +1684,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         y_all: torch.Tensor,
         i: int,
         steps: int,
-        prev_generated: Union[torch.Tensor, List[torch.Tensor]],
+        prev_generated: Union[torch.Tensor, list[torch.Tensor]],
         finished: bool = False,
-    ) -> Tuple[Union[torch.Tensor, List[torch.Tensor]], torch.Tensor]:
+    ) -> tuple[Union[torch.Tensor, list[torch.Tensor]], torch.Tensor]:
         _mel, out = self.process_little_chunk(
             conditioning=conditioning,
             reference_mel=reference_mel,
@@ -1917,9 +1694,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
             y_all=y_all,
             i=i,
             steps=steps,
-            prev_generated=(
-                prev_generated if isinstance(prev_generated, torch.Tensor) else None
-            ),
+            prev_generated=(prev_generated if isinstance(prev_generated, torch.Tensor) else None),
             finished=finished,
         )
         return _mel if _mel is not None else prev_generated, out

vllm_omni/model_executor/models/qwen2_old.py

@@ -1,9 +1,9 @@
-from typing import Iterable, Optional, Set, Tuple, Union
+from collections.abc import Iterable
+from typing import Optional, Union
 
 import torch
 from torch import nn
 from transformers import Qwen2Config
-
 from vllm.attention import Attention, AttentionType
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -11,23 +11,13 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import (
-    MergedColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-)
+from vllm.model_executor.layers.linear import MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.pooler import Pooler, PoolingType
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead,
-    VocabParallelEmbedding,
-)
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
-)
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead, VocabParallelEmbedding
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader, maybe_remap_kv_scale_name
 from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
@@ -48,7 +38,6 @@ logger = init_logger(__name__)
 
 
 class Qwen2MLP(nn.Module):
-
     def __init__(
         self,
         hidden_size: int,
@@ -73,10 +62,7 @@ class Qwen2MLP(nn.Module):
             prefix=f"{prefix}.down_proj",
         )
         if hidden_act != "silu":
-            raise ValueError(
-                f"Unsupported activation: {hidden_act}. "
-                "Only silu is supported for now."
-            )
+            raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.")
         self.act_fn = SiluAndMul()
 
     def forward(self, x):
@@ -87,7 +73,6 @@ class Qwen2MLP(nn.Module):
 
 
 class Qwen2Attention(nn.Module):
-
     def __init__(
         self,
         hidden_size: int,
@@ -98,7 +83,7 @@ class Qwen2Attention(nn.Module):
         head_dim: Optional[int] = None,
         cache_config: Optional[CacheConfig] = None,
         quant_config: Optional[QuantizationConfig] = None,
-        rope_scaling: Optional[Tuple] = None,
+        rope_scaling: Optional[tuple] = None,
         prefix: str = "",
         attn_type: str = AttentionType.DECODER,
     ) -> None:
@@ -173,7 +158,6 @@ class Qwen2Attention(nn.Module):
 
 
 class Qwen2DecoderLayer(nn.Module):
-
     def __init__(
         self,
         config: Qwen2Config,
@@ -217,16 +201,14 @@ class Qwen2DecoderLayer(nn.Module):
             prefix=f"{prefix}.mlp",
         )
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         residual: Optional[torch.Tensor],
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         # Self Attention
         if residual is None:
             residual = hidden_states
@@ -255,7 +237,6 @@ class Qwen2DecoderLayer(nn.Module):
     }
 )
 class Qwen2Model(nn.Module):
-
     def __init__(
         self,
         *,
@@ -270,27 +251,20 @@ class Qwen2Model(nn.Module):
         quant_config = vllm_config.quant_config
 
         # TODO (@robertgshaw2): see if this can be moved out
-        if cache_config.sliding_window is not None and hasattr(
-            config, "max_window_layers"
-        ):
+        if cache_config.sliding_window is not None and hasattr(config, "max_window_layers"):
             raise ValueError(
                 "Sliding window for some but all layers is not "
                 "supported. This model uses sliding window "
-                "but `max_window_layers` = {} is less than "
-                "`num_hidden_layers` = {}. Please open an issue "
-                "to discuss this feature.".format(
-                    config.max_window_layers,
-                    config.num_hidden_layers,
-                )
+                f"but `max_window_layers` = {config.max_window_layers} is less than "
+                f"`num_hidden_layers` = {config.num_hidden_layers}. Please open an issue "
+                "to discuss this feature."
             )
 
         self.config = config
         self.quant_config = quant_config
         self.vocab_size = config.vocab_size
 
-        if get_pp_group().is_first_rank or (
-            config.tie_word_embeddings and get_pp_group().is_last_rank
-        ):
+        if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank):
             self.embed_tokens = VocabParallelEmbedding(
                 config.vocab_size,
                 getattr(config, "embedding_size", config.hidden_size),
@@ -348,13 +322,11 @@ class Qwen2Model(nn.Module):
                 residual,
             )
         if not get_pp_group().is_last_rank:
-            return IntermediateTensors(
-                {"hidden_states": hidden_states, "residual": residual}
-            )
+            return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
-    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
@@ -364,19 +336,15 @@ class Qwen2Model(nn.Module):
             ("gate_up_proj", "up_proj", 1),
         ]
         params_dict = dict(self.named_parameters(remove_duplicate=False))
-        loaded_params: Set[str] = set()
+        loaded_params: set[str] = set()
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
-            if self.quant_config is not None and (
-                scale_name := self.quant_config.get_cache_scale(name)
-            ):
+            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
                 # Loading kv cache quantization scales
                 param = params_dict[scale_name]
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = (
-                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
-                )
+                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
                 weight_loader(param, loaded_weight)
                 loaded_params.add(scale_name)
                 continue
@@ -433,9 +401,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         self.lora_config = lora_config
 
         self.quant_config = quant_config
-        self.model = Qwen2Model(
-            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
-        )
+        self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
         if get_pp_group().is_last_rank:
             if config.tie_word_embeddings:
@@ -453,9 +419,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         self.logits_processor = LogitsProcessor(config.vocab_size)
         self.sampler = Sampler()
 
-        self.make_empty_intermediate_tensors = (
-            self.model.make_empty_intermediate_tensors
-        )
+        self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
 
     def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.get_input_embeddings(input_ids)
@@ -467,9 +431,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         intermediate_tensors: Optional[IntermediateTensors] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, IntermediateTensors]:
-        hidden_states = self.model(
-            input_ids, positions, intermediate_tensors, inputs_embeds
-        )
+        hidden_states = self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
         return hidden_states
 
     def compute_logits(
@@ -487,7 +449,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         next_tokens = self.sampler(logits, sampling_metadata)
         return next_tokens
 
-    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(
             self,
             skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
@@ -521,9 +483,7 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
         self.lora_config = lora_config
 
         self.quant_config = quant_config
-        self.model = Qwen2Model(
-            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
-        )
+        self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
         # TODO: Replace this model class with as_embedding_model(
         # Qwen2ForCausalLM) after changing the default pooling method
@@ -554,9 +514,7 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
     ) -> Optional[PoolerOutput]:
         return self._pooler(hidden_states, pooling_metadata)
 
-    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         weights = self.hf_to_vllm_mapper.apply(weights)
-        weights = (
-            (name, data) for name, data in weights if not name.startswith("lm_head.")
-        )
+        weights = ((name, data) for name, data in weights if not name.startswith("lm_head."))
         self.model.load_weights(weights)

vllm_omni/model_executor/models/registry.py

@@ -1,8 +1,4 @@
-from vllm.model_executor.models.registry import (
-    _VLLM_MODELS,
-    _LazyRegisteredModel,
-    _ModelRegistry,
-)
+from vllm.model_executor.models.registry import _VLLM_MODELS, _LazyRegisteredModel, _ModelRegistry
 
 _OMNI_MODELS = {
     "Qwen2_5OmniForConditionalGeneration": (

vllm_omni/model_executor/ode_solver/__init__.py

@@ -0,0 +1,15 @@
+from vllm_omni.model_executor.ode_solver.runger_kutta_4 import RungeKutta4ODESolver
+
+_ODE_SOLVERS = {
+    "RungeKutta4ODESolver": RungeKutta4ODESolver,
+}
+
+
+def get_ode_solver_class(class_name: str):
+    solver = _ODE_SOLVERS.get(class_name)
+    if solver is None:
+        raise ValueError(f"ODE solver class '{class_name}' not found.")
+    return solver
+
+
+__all__ = ["RungeKutta4ODESolver"]

vllm_omni/model_executor/ode_solver/runger_kutta_4.py

@@ -0,0 +1,89 @@
+from typing import Callable
+
+import torch
+
+
+class RungeKutta4ODESolver:
+    def __init__(self, function: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], initial_value: torch.Tensor):
+        self.function = function
+        self.initial_value = initial_value
+
+        self._one_third = 1 / 3
+        self._two_thirds = 2 / 3
+
+    def reset_ode_solver(self, function, initial_value):
+        self.function = function
+        self.initial_value = initial_value
+
+    def _rk4_step(
+        self,
+        function,
+        time_start,
+        time_step,
+        time_end,
+        value_start,
+        function_value_start=None,
+    ):
+        k1 = function_value_start if function_value_start is not None else function(time_start, value_start)
+        k2 = function(
+            time_start + time_step * self._one_third,
+            value_start + time_step * k1 * self._one_third,
+        )
+        k3 = function(
+            time_start + time_step * self._two_thirds,
+            value_start + time_step * (k2 - k1 * self._one_third),
+        )
+        k4 = function(time_end, value_start + time_step * (k1 - k2 + k3))
+        return (k1 + 3 * (k2 + k3) + k4) * time_step / 8
+
+    def _compute_step(self, function, time_start, time_step, time_end, value_start):
+        function_value_start = function(time_start, value_start)
+        return (
+            self._rk4_step(
+                function,
+                time_start,
+                time_step,
+                time_end,
+                value_start,
+                function_value_start=function_value_start,
+            ),
+            function_value_start,
+        )
+
+    def _linear_interpolation(self, time_start, time_end, value_start, value_end, time_point):
+        if time_point == time_start:
+            return value_start
+        if time_point == time_end:
+            return value_end
+        weight = (time_point - time_start) / (time_end - time_start)
+        return value_start + weight * (value_end - value_start)
+
+    def integrate(self, time_points: torch.Tensor) -> torch.Tensor:
+        solution = torch.empty(
+            len(time_points),
+            *self.initial_value.shape,
+            dtype=self.initial_value.dtype,
+            device=self.initial_value.device,
+        )
+        solution[0] = self.initial_value
+
+        current_index = 1
+        current_value = self.initial_value
+        for time_start, time_end in zip(time_points[:-1], time_points[1:]):
+            time_step = time_end - time_start
+            delta_value, _ = self._compute_step(self.function, time_start, time_step, time_end, current_value)
+            next_value = current_value + delta_value
+
+            while current_index < len(time_points) and time_end >= time_points[current_index]:
+                solution[current_index] = self._linear_interpolation(
+                    time_start,
+                    time_end,
+                    current_value,
+                    next_value,
+                    time_points[current_index],
+                )
+                current_index += 1
+
+            current_value = next_value
+
+        return solution

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -8,15 +8,24 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu.gpu_ar_worker.GPUARWorker
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # need to discuss
       trust_remote_code: true
       engine_output_type: latent  # change the param name,such as pooling_output
       enable_prefix_caching: false
+    is_comprehension: true
     final_output: true
     final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
   - stage_id: 1
     runtime:
       process: true
@@ -25,7 +34,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu.gpu_ar_worker.GPUARWorker
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -34,6 +43,15 @@ stage_args:
       engine_output_type: latent
     engine_input_source: [0]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
+    default_sampling_params:
+      temperature: 0.9
+      top_p: 0.8
+      top_k: 40
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+      stop_token_ids: [8294]
   - stage_id: 2
     runtime:
       process: true
@@ -42,7 +60,8 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu.gpu_diffusion_worker.GPUDiffusionWorker
+      ode_solver_class: RungeKutta4ODESolver
+      worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true
@@ -52,6 +71,14 @@ stage_args:
     engine_input_source: [1]
     final_output: true
     final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
 
 # Top-level runtime config (concise): default windows and stage edges
 runtime:

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -1,15 +1,16 @@
 from typing import Union
 
 import torch
-
 from vllm.inputs import TextPrompt
+
 from vllm_omni.inputs.data import OmniTokensPrompt
+
 TALKER_CODEC_PAD_TOKEN_ID = 8292
 TALKER_CODEC_START_TOKEN_ID = 8293
 TALKER_CODEC_END_TOKEN_ID = 8294
-def thinker2talker(
-    stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None
-):
+
+
+def thinker2talker(stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None):
     if not engine_input_source:
         raise ValueError("engine_input_source cannot be empty")
     source_stage_id = engine_input_source[0]
@@ -19,9 +20,10 @@ def thinker2talker(
         raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
     thinker_outputs = stage_list[source_stage_id].engine_outputs
     talker_inputs = []
+    if not isinstance(prompt, list):
+        prompt = [prompt]
     multi_modal_data = {
-        thinker_output.request_id: prompt.get("multi_modal_data", None)
-        for thinker_output, prompt in zip(thinker_outputs, prompt)
+        thinker_output.request_id: p.get("multi_modal_data", None) for thinker_output, p in zip(thinker_outputs, prompt)
     }
 
     for i, thinker_output in enumerate(thinker_outputs):
@@ -29,28 +31,24 @@ def thinker2talker(
         prompt_token_ids = thinker_output.prompt_token_ids
         thinker_output_ids = output.token_ids
         prompt_token_ids_len = len(prompt_token_ids)
-        thinker_hidden_states = (
-            output.multimodal_output["latent"].clone().detach().cuda()
-        )
+        latent = output.multimodal_output["latent"]
+        thinker_hidden_states = latent.clone().detach().to(latent.device)
+        additional_information = {
+            "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(torch.float32),
+            "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(torch.float32),
+            "prompt_token_ids": prompt_token_ids,
+            "thinker_output_token_ids": thinker_output_ids,
+            "thinker_result_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape),
+            "prompt_embeds_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape),
+        }
         talker_inputs.append(
             OmniTokensPrompt(
-                prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID] + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids)) + [TALKER_CODEC_END_TOKEN_ID],
-                additional_information={
-                    "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(
-                        torch.float32
-                    ),
-                    "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(
-                        torch.float32
-                    ),
-                    "prompt_token_ids": prompt_token_ids,
-                    "thinker_output_token_ids": thinker_output_ids,
-                    "thinker_result_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape),
-                    "prompt_embeds_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape),
-                },
+                prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID]
+                + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids))
+                + [TALKER_CODEC_END_TOKEN_ID],
+                additional_information=additional_information,
                 multi_modal_data=(
-                    multi_modal_data[thinker_output.request_id]
-                    if multi_modal_data is not None
-                    else None
+                    multi_modal_data[thinker_output.request_id] if multi_modal_data is not None else None
                 ),
                 mm_processor_kwargs=None,
             )

vllm_omni/outputs.py

@@ -2,7 +2,6 @@ from dataclasses import dataclass
 from typing import Optional
 
 import torch
-
 from vllm.outputs import RequestOutput
 from vllm.v1.outputs import ModelRunnerOutput
 

vllm_omni/patch.py

@@ -1,11 +1,10 @@
 import sys
 
 from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
-from vllm.model_executor.layers.rotary_embedding import (
-    MRotaryEmbedding as _OriginalMRotaryEmbedding,
-)
+from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding as _OriginalMRotaryEmbedding
 from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest
 from vllm.v1.request import Request as _OriginalRequest
+
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
@@ -14,15 +13,9 @@ from vllm_omni.request import OmniRequest
 for module_name, module in sys.modules.items():
     if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt:
         module.TokensPrompt = OmniTokensPrompt
-    if (
-        hasattr(module, "MRotaryEmbedding")
-        and module.MRotaryEmbedding == _OriginalMRotaryEmbedding
-    ):
+    if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding:
         module.MRotaryEmbedding = MRotaryEmbedding
     if hasattr(module, "Request") and module.Request == _OriginalRequest:
         module.Request = OmniRequest
-    if (
-        hasattr(module, "EngineCoreRequest")
-        and module.EngineCoreRequest == _OriginalEngineCoreRequest
-    ):
+    if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest:
         module.EngineCoreRequest = OmniEngineCoreRequest

vllm_omni/request.py

@@ -6,11 +6,7 @@ from vllm.v1.structured_output.request import StructuredOutputRequest
 if TYPE_CHECKING:
     from vllm.v1.core.kv_cache_utils import BlockHash
 
-from vllm_omni.engine import (
-    AdditionalInformationPayload,
-    OmniEngineCoreRequest,
-    PromptEmbedsPayload,
-)
+from vllm_omni.engine import AdditionalInformationPayload, OmniEngineCoreRequest, PromptEmbedsPayload
 
 
 class OmniRequest(Request):
@@ -25,9 +21,7 @@ class OmniRequest(Request):
         # Serialized prompt embeddings payload (optional)
         self.prompt_embeds: Optional[PromptEmbedsPayload] = prompt_embeds
         # Serialized additional information payload (optional)
-        self.additional_information: Optional[AdditionalInformationPayload] = (
-            additional_information
-        )
+        self.additional_information: Optional[AdditionalInformationPayload] = additional_information
 
     @classmethod
     def from_engine_core_request(
@@ -35,7 +29,6 @@ class OmniRequest(Request):
         request: OmniEngineCoreRequest,
         block_hasher: Optional[Callable[["Request"], list["BlockHash"]]],
     ) -> "Request":
-
         return cls(
             request_id=request.request_id,
             client_index=request.client_index,
@@ -47,9 +40,7 @@ class OmniRequest(Request):
             arrival_time=request.arrival_time,
             lora_request=request.lora_request,
             structured_output_request=(
-                StructuredOutputRequest(sampling_params=request.sampling_params)
-                if request.sampling_params
-                else None
+                StructuredOutputRequest(sampling_params=request.sampling_params) if request.sampling_params else None
             ),
             cache_salt=request.cache_salt,
             priority=request.priority,

vllm_omni/version.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+try:
+    from ._version import __version__, __version_tuple__
+except Exception as e:
+    import warnings
+
+    warnings.warn(f"Failed to read commit hash:\n{e}", RuntimeWarning, stacklevel=2)
+
+    __version__ = "dev"
+    __version_tuple__ = (0, 0, __version__)

vllm_omni/worker/cpu/cpu_ar_model_runner.py

@@ -1,649 +0,0 @@
-"""AR CPU Model Runner for vLLM-omni.
-
-Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
-and also outputs sampled tokens.
-"""
-
-from __future__ import annotations
-
-from typing import Any, Optional, Union
-
-import numpy as np
-import torch
-
-from vllm.forward_context import BatchDescriptor
-from vllm.logger import init_logger
-from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.outputs import AsyncModelRunnerOutput
-from vllm.v1.structured_output.utils import apply_grammar_bitmask
-from vllm.v1.utils import record_function_or_nullcontext
-from vllm.v1.worker.gpu_model_runner import (
-    EMPTY_MODEL_RUNNER_OUTPUT,
-    AsyncGPUModelRunnerOutput,
-    IntermediateTensors,
-    get_pp_group,
-    get_tp_group,
-    has_kv_transfer_group,
-    set_forward_context,
-)
-from vllm.v1.worker.ubatch_utils import UBatchSlices
-from vllm.v1.worker.utils import is_residual_scattered_for_sp
-from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
-from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.cpu.cpu_model_runner import OmniCPUModelRunner
-
-logger = init_logger(__name__)
-
-
-class CPUARModelRunner(OmniCPUModelRunner):
-    """Autoregressive GPU model runner that returns hidden states per request.
-
-    This runner follows the same preparation and forward path as GPUModelRunner
-    (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs),
-    and additionally performs lightweight sampling so that sampled tokens are
-    available in outputs. Hidden representations are taken at the same indices
-    that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`).
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
-
-    def _preprocess(
-        self,
-        scheduler_output: "SchedulerOutput",
-        num_scheduled_tokens_np: np.ndarray,
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-        ubatch_slices: Optional[UBatchSlices] = None,
-        num_tokens_after_padding: Optional[torch.Tensor] = None,
-    ) -> tuple[
-        int,
-        int,
-        Optional[torch.Tensor],
-        Optional[torch.Tensor],
-        Optional[torch.Tensor],
-        torch.Tensor,
-        Optional[IntermediateTensors],
-        dict[str, Any],
-        Optional[dict[str, dict]],
-    ]:
-
-        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        if ubatch_slices:
-            assert num_tokens_after_padding is not None
-            num_input_tokens = int(num_tokens_after_padding[0].item() * 2)
-            self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
-        elif ubatch_slices is None:
-            num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)
-            num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
-            num_input_tokens += num_pad
-
-        # _prepare_inputs may reorder the batch, so we must gather multi
-        # modal outputs after that to ensure the correct order
-        per_req_additional_information: Optional[dict[str, dict]] = None
-        if (
-            self.supports_mm_inputs
-            and get_pp_group().is_first_rank
-            and not self.model_config.is_encoder_decoder
-        ):
-            # Run the multimodal encoder if any.
-            self._execute_mm_encoder(scheduler_output)
-            mm_embeds = self._gather_mm_embeddings(scheduler_output)
-
-            # NOTE(woosuk): To unify token ids and soft tokens (vision
-            # embeddings), we always use embeddings (rather than token ids)
-            # as input to the multimodal model, even when the input is text.
-            inputs_embeds_scheduled = self.model.get_input_embeddings(
-                input_ids=self.input_ids.gpu[:num_scheduled_tokens],
-                multimodal_embeddings=mm_embeds or None,
-            )
-
-            # TODO(woosuk): Avoid the copy. Optimize.
-            self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
-
-            # Reset per-step additional information collector (deprecated concat path)
-            if hasattr(self, "_forward_additional_information"):
-                self._forward_additional_information = None
-            # New: per-request additional information for this step
-            per_req_additional_information = {}
-
-            # Overlay custom prompt_embeds per request for the prompt portion;
-            # collect additional_information (tensor/list) for prefill portion only
-            for req_index, req_id in enumerate(self.input_batch.req_ids):
-                req_state = self.requests[req_id]
-                pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
-                addi_cpu = getattr(req_state, "additional_information_cpu", None)
-                num_computed_tokens = int(
-                    self.input_batch.num_computed_tokens_cpu[req_index]
-                )
-                prompt_len = len(req_state.prompt_token_ids)
-                prompt_remaining = max(0, prompt_len - num_computed_tokens)
-                sched_tokens = int(num_scheduled_tokens_np[req_index])
-                overlay_len = min(sched_tokens, prompt_remaining)
-                if overlay_len <= 0:
-                    continue
-                if pe_cpu is not None:
-                    src = pe_cpu[
-                        num_computed_tokens : num_computed_tokens + overlay_len
-                    ].to(dtype=self.dtype, device=self.device, non_blocking=True)
-                    start_offset = int(self.query_start_loc.cpu[req_index])
-                    self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(
-                        src
-                    )
-                # Build per-request additional information (no cross-request concat)
-                if addi_cpu is not None and isinstance(addi_cpu, dict):
-                    req_info: dict[str, object] = {}
-                    for k, v in addi_cpu.items():
-                        if isinstance(v, torch.Tensor):
-                            # For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor
-                            if overlay_len > 0:
-                                try:
-                                    seg = v[
-                                        num_computed_tokens : num_computed_tokens + overlay_len
-                                    ].detach().to("cpu").contiguous()
-                                except Exception:
-                                    seg = v.detach().to("cpu").contiguous()
-                                req_info[k] = seg
-                            else:
-                                req_info[k] = v.detach().to("cpu").contiguous()
-                        elif isinstance(v, list):
-                            req_info[k] = v
-                        else:
-                            req_info[k] = v
-                    per_req_additional_information[req_id] = req_info
-
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = {
-                **self._init_model_kwargs(num_scheduled_tokens),
-                **self._extract_mm_kwargs(scheduler_output),
-            }
-        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
-            # Get the input embeddings for the tokens that are not input embeds,
-            # then put them into the appropriate positions.
-            # TODO(qthequartermasterman): Since even when prompt embeds are
-            # enabled, (a) not all requests will use prompt embeds, and (b)
-            # after the initial prompt is processed, the rest of the generated
-            # tokens will be token ids, it is not desirable to have the
-            # embedding layer outside of the CUDA graph all the time. The v0
-            # engine avoids this by "double compiling" the CUDA graph, once
-            # with input_ids and again with inputs_embeds, for all num_tokens.
-            # If a batch only has token ids, then including the embedding layer
-            # in the CUDA graph will be more performant (like in the else case
-            # below).
-            token_ids_idx = (
-                self.is_token_ids.gpu[:num_scheduled_tokens]
-                .nonzero(as_tuple=False)
-                .squeeze(1)
-            )
-            # Some tokens ids may need to become embeds
-            if token_ids_idx.numel() > 0:
-                token_ids = self.input_ids.gpu[token_ids_idx]
-                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
-                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
-
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = self._init_model_kwargs(num_input_tokens)
-            input_ids = None
-        else:
-            # For text-only models, we use token ids as input.
-            # While it is possible to use embeddings as input just like the
-            # multimodal models, it is not desirable for performance since
-            # then the embedding layer is not included in the CUDA graph.
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = None
-            model_kwargs = self._init_model_kwargs(num_input_tokens)
-        if self.uses_mrope:
-            positions = self.mrope_positions.gpu[:, :num_input_tokens]
-        else:
-            positions = self.positions.gpu[:num_input_tokens]
-
-        if get_pp_group().is_first_rank:
-            intermediate_tensors = None
-        else:
-            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True
-            )
-
-        if (
-            self.model_config.is_encoder_decoder
-            and scheduler_output.scheduled_encoder_inputs
-        ):
-            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
-            model_kwargs.update(encoder_inputs)
-
-        return (
-            num_scheduled_tokens,
-            num_input_tokens,
-            num_tokens_after_padding,
-            input_ids,
-            inputs_embeds,
-            positions,
-            intermediate_tensors,
-            model_kwargs,
-            per_req_additional_information,
-        )
-
-    @torch.inference_mode()
-    def execute_model(
-        self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-    ) -> Union[OmniModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:
-        with record_function_or_nullcontext("Preprocess"):
-            with self.synchronize_input_prep():
-                # Update persistent batch states.
-                self._update_states(scheduler_output)
-
-                # Decode per-request prompt_embeds / additional_hidden_states payloads
-                # (if present) into CPU tensors
-                try:
-                    new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
-                    if new_reqs:
-                        for nr in new_reqs:
-                            req_id = getattr(nr, "req_id", None) or getattr(
-                                nr, "request_id", None
-                            )
-                            if req_id is None:
-                                continue
-                            # prompt_embeds
-                            payload_pe = getattr(nr, "prompt_embeds", None)
-                            if payload_pe is not None:
-                                if isinstance(payload_pe, torch.Tensor):
-                                    pe_cpu = payload_pe.detach().to("cpu").contiguous()
-                                elif isinstance(payload_pe, PromptEmbedsPayload):
-                                    dt = np.dtype(
-                                        getattr(payload_pe, "dtype", "float32")
-                                    )
-                                    arr = np.frombuffer(payload_pe.data, dtype=dt)
-                                    arr = arr.reshape(payload_pe.shape)
-                                    pe_cpu = torch.from_numpy(arr)
-                                else:
-                                    pe_cpu = None
-                                if pe_cpu is not None and req_id in self.requests:
-                                    setattr(
-                                        self.requests[req_id],
-                                        "prompt_embeds_cpu",
-                                        pe_cpu,
-                                    )
-                            # additional_information
-                            payload_info = getattr(nr, "additional_information", None)
-                            if payload_info is not None:
-                                info_dict = {}
-                                if isinstance(payload_info, dict):
-                                    # Already decoded
-                                    info_dict = payload_info
-                                elif isinstance(
-                                    payload_info, AdditionalInformationPayload
-                                ):
-                                    for k, entry in payload_info.entries.items():
-                                        if entry.tensor_data is not None:
-                                            dt = np.dtype(
-                                                getattr(
-                                                    entry, "tensor_dtype", "float32"
-                                                )
-                                            )
-                                            arr = np.frombuffer(
-                                                entry.tensor_data, dtype=dt
-                                            )
-                                            arr = arr.reshape(entry.tensor_shape)
-                                            info_dict[k] = torch.from_numpy(arr)
-                                        else:
-                                            info_dict[k] = entry.list_data
-                                if info_dict and req_id in self.requests:
-                                    setattr(
-                                        self.requests[req_id],
-                                        "additional_information_cpu",
-                                        info_dict,
-                                    )
-                except Exception as e:
-                    logger.error(
-                        f"Error decoding prompt_embeds / additional_information: {e}"
-                    )
-                    pass
-
-                if not scheduler_output.total_num_scheduled_tokens:
-                    if not has_kv_transfer_group():
-                        # Return empty ModelRunnerOutput if no work to do.
-                        return EMPTY_MODEL_RUNNER_OUTPUT
-                    return self.kv_connector_no_forward(
-                        scheduler_output, self.vllm_config
-                    )
-                if self.cache_config.kv_sharing_fast_prefill:
-                    assert not self.input_batch.num_prompt_logprobs, (
-                        "--kv-sharing-fast-prefill produces incorrect "
-                        "logprobs for prompt tokens, tokens, please disable "
-                        "it when the requests need prompt logprobs"
-                    )
-
-                # Prepare the decoder inputs.
-                (
-                    attn_metadata,
-                    logits_indices,
-                    spec_decode_metadata,
-                    num_scheduled_tokens_np,
-                    spec_decode_common_attn_metadata,
-                    max_query_len,
-                    ubatch_slices,
-                    num_tokens_after_padding,
-                ) = self._prepare_inputs(scheduler_output)
-
-            (
-                num_scheduled_tokens,
-                num_input_tokens,
-                num_tokens_across_dp,
-                input_ids,
-                inputs_embeds,
-                positions,
-                intermediate_tensors,
-                model_kwargs,
-                per_req_additional_information,
-            ) = self._preprocess(
-                scheduler_output,
-                num_scheduled_tokens_np,
-                intermediate_tensors,
-                ubatch_slices,
-                num_tokens_after_padding,
-            )
-
-            uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
-                num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
-            )
-            batch_descriptor = BatchDescriptor(
-                num_tokens=num_input_tokens, uniform_decode=uniform_decode
-            )
-            cudagraph_runtime_mode, batch_descriptor = (
-                self.cudagraph_dispatcher.dispatch(batch_descriptor)
-            )
-
-        # This is currently to get around the assert in the DPMetadata
-        # where it wants `num_tokens_across_dp` to align with `num_tokens`
-        if ubatch_slices is not None:
-            num_input_tokens = ubatch_slices[0].num_tokens
-
-        # Run the model.
-        # Use persistent buffers for CUDA graphs.
-        with (
-            set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_input_tokens,
-                num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_runtime_mode,
-                batch_descriptor=batch_descriptor,
-                ubatch_slices=ubatch_slices,
-            ),
-            record_function_or_nullcontext("Forward"),
-            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
-        ):
-
-            model_kwargs_extra = {}
-            # Pass per-request additional information map for this step (no concat)
-            if per_req_additional_information:
-                model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
-            # Always pass per-request runtime additional_information (persisted in request state)
-            try:
-                per_req_runtime_info = []
-                for req_id in self.input_batch.req_ids:
-                    req_state = self.requests.get(req_id)
-                    info = (
-                        getattr(req_state, "additional_information_cpu", None)
-                        if req_state is not None
-                        else None
-                    )
-                    per_req_runtime_info.append(info if isinstance(info, dict) else {})
-                model_kwargs_extra["runtime_additional_information"] = (
-                    per_req_runtime_info
-                )
-                model_kwargs_extra["request_ids"] = self.input_batch.req_ids
-                # Pass each request's token span within the flattened sequence for this step, enabling the model to map decode/prefill by request
-                req_token_spans = []
-                for req_index in range(len(self.input_batch.req_ids)):
-                    start_offset = int(self.query_start_loc.cpu[req_index])
-                    sched_tokens = int(num_scheduled_tokens_np[req_index])
-                    req_token_spans.append((start_offset, start_offset + sched_tokens))
-                model_kwargs_extra["request_token_spans"] = req_token_spans
-            except Exception:
-                pass
-            model_output = self.model(
-                input_ids=input_ids,
-                positions=positions,
-                intermediate_tensors=intermediate_tensors,
-                inputs_embeds=inputs_embeds,
-                **model_kwargs,
-                sampling_metadata=self.input_batch.sampling_metadata,
-                logits_index=logits_indices,
-                sampler=self.sampler,
-                **model_kwargs_extra,
-            )
-
-        with record_function_or_nullcontext("Postprocess"):
-            if self.use_aux_hidden_state_outputs:
-                # True when EAGLE 3 is used.
-                hidden_states, aux_hidden_states = model_output
-            else:
-                # Common case.
-                hidden_states = model_output
-                aux_hidden_states = None
-
-            hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
-                hidden_states
-            )
-            # The model side may return per-request additional_information updates (model-agnostic channel).
-            # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
-            # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
-            try:
-                if (
-                    isinstance(multimodal_outputs, dict)
-                    and (
-                        "additional_information_update" in multimodal_outputs
-                        or "additional_information_update_by_req_id" in multimodal_outputs
-                    )
-                ):
-                    # Option A: list[dict] in batch order
-                    updates_list = multimodal_outputs.get(
-                        "additional_information_update"
-                    )
-                    if isinstance(updates_list, list):
-                        for idx, upd in enumerate(updates_list):
-                            if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
-                                continue
-                            req_id = self.input_batch.req_ids[idx]
-                            self._merge_additional_information_update(req_id, upd)
-                    # Option B: dict[str, dict] keyed by req_id
-                    updates_map = multimodal_outputs.get(
-                        "additional_information_update_by_req_id"
-                    )
-                    if isinstance(updates_map, dict):
-                        for req_id, upd in updates_map.items():
-                            if not isinstance(upd, dict):
-                                continue
-                            if req_id not in self.requests:
-                                continue
-                            self._merge_additional_information_update(req_id, upd)
-            except Exception as e:
-                logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}")
-            if not self.broadcast_pp_output:
-
-                # Common case.
-                if not get_pp_group().is_last_rank:
-                    # Return the intermediate tensors.
-                    assert isinstance(hidden_states, IntermediateTensors)
-                    hidden_states.kv_connector_output = kv_connector_output
-                    return hidden_states
-
-                if self.is_pooling_model:
-                    # Return the pooling output.
-                    output = self._pool(
-                        hidden_states, num_scheduled_tokens, num_scheduled_tokens_np
-                    )
-                    output.kv_connector_output = kv_connector_output
-                    return output
-
-                sample_hidden_states = hidden_states[logits_indices]
-                logits = self.model.compute_logits(sample_hidden_states)
-            else:
-                # Rare case.
-                assert not self.is_pooling_model
-
-                if not get_pp_group().is_last_rank:
-                    all_gather_tensors = {
-                        "residual": not is_residual_scattered_for_sp(
-                            self.vllm_config, num_input_tokens
-                        )
-                    }
-                    get_pp_group().send_tensor_dict(
-                        hidden_states.tensors,
-                        all_gather_group=get_tp_group(),
-                        all_gather_tensors=all_gather_tensors,
-                    )
-                    logits = None
-                else:
-                    sample_hidden_states = hidden_states[logits_indices]
-                    logits = self.model.compute_logits(sample_hidden_states)
-
-                model_output_broadcast_data = {}
-                if logits is not None:
-                    model_output_broadcast_data["logits"] = logits.contiguous()
-
-                model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
-                    model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
-                )
-                assert model_output_broadcast_data is not None
-                logits = model_output_broadcast_data["logits"]
-
-            # Apply structured output bitmasks if present
-            if scheduler_output.grammar_bitmask is not None:
-                apply_grammar_bitmask(
-                    scheduler_output, self.input_batch, logits, self.device
-                )
-
-        with record_function_or_nullcontext("Sample"):
-            sampler_output = self._sample(logits, spec_decode_metadata)
-
-        def propose_draft_token_ids(sampled_token_ids):
-            assert spec_decode_common_attn_metadata is not None
-            with record_function_or_nullcontext("Draft"):
-                self._draft_token_ids = self.propose_draft_token_ids(
-                    scheduler_output,
-                    sampled_token_ids,
-                    self.input_batch.sampling_metadata,
-                    hidden_states,
-                    sample_hidden_states,
-                    aux_hidden_states,
-                    spec_decode_metadata,
-                    spec_decode_common_attn_metadata,
-                )
-
-        use_padded_batch_for_eagle = (
-            self.speculative_config
-            and self.speculative_config.use_eagle()
-            and not self.speculative_config.disable_padded_drafter_batch
-        )
-        effective_drafter_max_model_len = self.max_model_len
-        if effective_drafter_max_model_len is None:
-            effective_drafter_max_model_len = self.model_config.max_model_len
-        if (
-            self.speculative_config
-            and self.speculative_config.draft_model_config is not None
-            and self.speculative_config.draft_model_config.max_model_len is not None
-        ):
-            effective_drafter_max_model_len = (
-                self.speculative_config.draft_model_config.max_model_len
-            )
-        input_fits_in_drafter = spec_decode_common_attn_metadata and (
-            spec_decode_common_attn_metadata.seq_lens.max()
-            + self.speculative_config.num_speculative_tokens
-            <= effective_drafter_max_model_len
-        )
-        if use_padded_batch_for_eagle and input_fits_in_drafter:
-            # EAGLE speculative decoding can use the GPU sampled tokens
-            # as inputs, and does not need to wait for bookkeeping to finish.
-            propose_draft_token_ids(sampler_output.sampled_token_ids)
-
-        with record_function_or_nullcontext("Bookkeep"):
-            (
-                num_nans_in_logits,
-                logprobs_lists,
-                valid_sampled_token_ids,
-                prompt_logprobs_dict,
-                req_ids_output_copy,
-                req_id_to_index_output_copy,
-                invalid_req_indices,
-            ) = self._bookkeeping_sync(
-                scheduler_output,
-                sampler_output,
-                logits,
-                hidden_states,
-                num_scheduled_tokens,
-            )
-
-        if (
-            self.speculative_config
-            and not use_padded_batch_for_eagle
-            and input_fits_in_drafter
-        ):
-            # ngram and other speculative decoding methods use the sampled
-            # tokens on the CPU, so they are run after bookkeeping.
-            propose_draft_token_ids(valid_sampled_token_ids)
-
-        with record_function_or_nullcontext("EPLB"):
-            self.eplb_step()
-
-        # Convert to per-request tensors on CPU
-        hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
-        pooler_output: list[Optional[torch.Tensor]] = []
-        prev_logits_index = 0
-        for logits_index in logits_indices:
-            pooler_output.append(
-                hidden_states_cpu[prev_logits_index : logits_index + 1]
-            )
-            prev_logits_index = logits_index + 1
-
-        output = OmniModelRunnerOutput(
-            req_ids=req_ids_output_copy,
-            req_id_to_index=req_id_to_index_output_copy,
-            sampled_token_ids=valid_sampled_token_ids,
-            logprobs=logprobs_lists,
-            prompt_logprobs_dict=prompt_logprobs_dict,
-            pooler_output=(
-                pooler_output
-                if self.vllm_config.model_config.engine_output_type != "text"
-                else None
-            ),
-            kv_connector_output=kv_connector_output,
-            num_nans_in_logits=num_nans_in_logits,
-        )
-
-        if not self.use_async_scheduling:
-            return output
-
-        return AsyncGPUModelRunnerOutput(
-            model_runner_output=output,
-            sampled_token_ids=sampler_output.sampled_token_ids,
-            invalid_req_indices=invalid_req_indices,
-            async_output_copy_stream=self.async_output_copy_stream,
-        )
-
-    def _merge_additional_information_update(self, req_id: str, upd: dict) -> None:
-        req_state = self.requests.get(req_id)
-        if req_state is None:
-            return
-        existing = getattr(req_state, "additional_information_cpu", {})
-        if not isinstance(existing, dict):
-            existing = {}
-        merged = dict(existing)
-        for k, v in upd.items():
-            if isinstance(v, torch.Tensor):
-                merged[k] = v.detach().to("cpu").contiguous()
-            elif isinstance(v, list):
-                new_list = []
-                for item in v:
-                    if isinstance(item, torch.Tensor):
-                        new_list.append(item.detach().to("cpu").contiguous())
-                    else:
-                        new_list.append(item)
-                merged[k] = new_list
-            else:
-                merged[k] = v
-        setattr(req_state, "additional_information_cpu", merged)

vllm_omni/worker/cpu/cpu_ar_worker.py

@@ -1,63 +0,0 @@
-import os
-import platform
-
-from vllm import envs
-import torch
-from vllm.logger import init_logger
-from vllm.model_executor import set_random_seed
-from vllm.platforms import current_platform, CpuArchEnum
-
-
-from vllm.v1.worker.cpu_worker import CPUWorker
-from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-
-from vllm_omni.worker.cpu.cpu_ar_model_runner import CPUARModelRunner
-
-logger = init_logger(__name__)
-
-
-class CPUARWorker(CPUWorker):
-    def init_device(self):
-        # Setup OpenMP threads affinity.
-        omp_cpuids = envs.VLLM_CPU_OMP_THREADS_BIND
-        if omp_cpuids == "auto" and platform.system() == "Linux":
-            cpu_arch = current_platform.get_cpu_architecture()
-            if cpu_arch in (CpuArchEnum.POWERPC, CpuArchEnum.S390X):
-                # For S390X/POWERPC SMT-8/4/2
-                self.local_omp_cpuid = self._get_autobind_cpu_ids(
-                    lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4])
-            elif current_platform.get_cpu_architecture() == CpuArchEnum.X86:
-                # For x86 SMT-2, use 1 CPU per core
-                self.local_omp_cpuid = self._get_autobind_cpu_ids(
-                    lambda cpus: cpus[-1:])
-            else:
-                self.local_omp_cpuid = "all"
-        else:
-            local_dp_rank = self.parallel_config.data_parallel_rank_local
-            omp_cpuids = omp_cpuids.split("|")
-            if local_dp_rank is not None:
-                world_size = self.parallel_config.world_size
-                omp_cpuids = omp_cpuids[local_dp_rank *
-                                        world_size:(local_dp_rank + 1) *
-                                                   world_size]
-            self.local_omp_cpuid = omp_cpuids[self.rank]
-
-        if self.local_omp_cpuid != "all":
-            ret = torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)
-            if ret:
-                logger.info(ret)
-
-        # Note: unique identifier for creating allreduce shared memory
-        os.environ["VLLM_DIST_IDENT"] = self.distributed_init_method.split(
-            ":")[-1]
-        # Initialize the distributed environment.
-        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)
-
-        # Construct the model runner
-        self.model_runner: CPUARModelRunner = CPUARModelRunner(
-            self.vllm_config, torch.device("cpu"))

vllm_omni/worker/cpu/cpu_diffusion_model_runner.py

@@ -1,620 +0,0 @@
-from __future__ import annotations
-
-import gc
-import logging
-from typing import Any, List, Optional, Union
-
-import numpy as np
-import torch
-
-from vllm.config import CUDAGraphMode
-from vllm.forward_context import BatchDescriptor
-from vllm.multimodal.inputs import MultiModalKwargs
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
-from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.spec_decode.eagle import EagleProposer
-from vllm.v1.utils import record_function_or_nullcontext
-from vllm.v1.worker.gpu_model_runner import (
-    EMPTY_MODEL_RUNNER_OUTPUT,
-    AsyncGPUModelRunnerOutput,
-    IntermediateTensors,
-    PerLayerAttnMetadata,
-    get_pp_group,
-    set_forward_context,
-)
-from vllm.v1.worker.ubatch_utils import UBatchSlices
-from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
-from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.cpu.cpu_model_runner import OmniCPUModelRunner
-
-logger = logging.getLogger(__name__)
-
-
-class CPUDiffusionModelRunner(OmniCPUModelRunner):
-    """Diffusion model runner for vLLM-omni (non-autoregressive).
-
-    - Reuses CPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
-    - Does not compute logits or perform token sampling.
-    - Executes diffusion process and returns tensors via `pooler_output`.
-    """
-
-    def _preprocess(
-        self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-        ubatch_slices: Optional[UBatchSlices] = None,
-        num_tokens_after_padding: Optional[torch.Tensor] = None,
-    ) -> tuple[
-        int,
-        int,
-        Optional[torch.Tensor],
-        Optional[torch.Tensor],
-        Optional[torch.Tensor],
-        torch.Tensor,
-        Optional[IntermediateTensors],
-        dict[str, Any],
-    ]:
-
-        # Input token count for this iteration (not used by diffusion, but
-        # retained to keep DP padding/ordering consistent)
-        num_input_tokens = scheduler_output.total_num_scheduled_tokens
-        num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
-        num_input_tokens += num_pad
-
-        # _prepare_inputs may reorder the batch, so we must gather multi
-        # modal outputs after that to ensure the correct order
-        if (
-            self.supports_mm_inputs
-            and get_pp_group().is_first_rank
-            and not self.model_config.is_encoder_decoder
-        ):
-            # Run the multimodal encoder if any.
-            self._execute_mm_encoder(scheduler_output)
-            mm_embeds = self._gather_mm_embeddings(scheduler_output)
-
-            # NOTE(woosuk): To unify token ids and soft tokens (vision
-            # embeddings), we always use embeddings (rather than token ids)
-            # as input to the multimodal model, even when the input is text.
-            inputs_embeds_scheduled = self.model.get_input_embeddings(
-                input_ids=self.input_ids.gpu[:num_input_tokens],
-                multimodal_embeddings=mm_embeds or None,
-            )
-
-            # TODO(woosuk): Avoid the copy. Optimize.
-            self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled)
-
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = {
-                **self._init_model_kwargs(num_input_tokens),
-                **self._extract_mm_kwargs(scheduler_output),
-            }
-        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
-            # Get the input embeddings for the tokens that are not input embeds,
-            # then put them into the appropriate positions.
-            # TODO(qthequartermasterman): Since even when prompt embeds are
-            # enabled, (a) not all requests will use prompt embeds, and (b)
-            # after the initial prompt is processed, the rest of the generated
-            # tokens will be token ids, it is not desirable to have the
-            # embedding layer outside of the CUDA graph all the time. The v0
-            # engine avoids this by "double compiling" the CUDA graph, once
-            # with input_ids and again with inputs_embeds, for all num_tokens.
-            # If a batch only has token ids, then including the embedding layer
-            # in the CUDA graph will be more performant (like in the else case
-            # below).
-            token_ids_idx = (
-                self.is_token_ids.gpu[:num_input_tokens]
-                .nonzero(as_tuple=False)
-                .squeeze(1)
-            )
-            # Some tokens ids may need to become embeds
-            if token_ids_idx.numel() > 0:
-                token_ids = self.input_ids.gpu[token_ids_idx]
-                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
-                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
-
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = self._init_model_kwargs(num_input_tokens)
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-        else:
-            # For text-only models, we use token ids as input.
-            # While it is possible to use embeddings as input just like the
-            # multimodal models, it is not desirable for performance since
-            # then the embedding layer is not included in the CUDA graph.
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = None
-            model_kwargs = self._init_model_kwargs(num_input_tokens)
-        if self.uses_mrope:
-            positions = self.mrope_positions.gpu[:, :num_input_tokens]
-        else:
-            positions = self.positions.gpu[:num_input_tokens]
-
-        if get_pp_group().is_first_rank:
-            intermediate_tensors = None
-        else:
-            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True
-            )
-
-        if (
-            self.model_config.is_encoder_decoder
-            and scheduler_output.scheduled_encoder_inputs
-        ):
-            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
-            model_kwargs.update(encoder_inputs)
-
-        return (
-            num_input_tokens,
-            num_input_tokens,
-            num_tokens_after_padding,
-            input_ids,
-            inputs_embeds,
-            positions,
-            intermediate_tensors,
-            model_kwargs,
-        )
-
-    @torch.inference_mode()
-    def execute_model(
-        self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-    ) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
-        with record_function_or_nullcontext("Preprocess"):
-            with self.synchronize_input_prep():
-                # Update persistent batch states.
-                self._update_states(scheduler_output)
-                if not scheduler_output.total_num_scheduled_tokens:
-                    return EMPTY_MODEL_RUNNER_OUTPUT
-
-                # Prepare the decoder inputs.
-                (
-                    attn_metadata,
-                    logits_indices,
-                    spec_decode_metadata,
-                    num_scheduled_tokens_np,
-                    spec_decode_common_attn_metadata,
-                    max_query_len,
-                    ubatch_slices,
-                    num_tokens_after_padding,
-                ) = self._prepare_inputs(scheduler_output)
-
-            # Input token count for this iteration (not used by diffusion, but
-            # retained to keep DP padding/ordering consistent)
-            (
-                num_scheduled_tokens,
-                num_input_tokens,
-                num_tokens_across_dp,
-                input_ids,
-                inputs_embeds,
-                positions,
-                intermediate_tensors,
-                model_kwargs,
-            ) = self._preprocess(
-                scheduler_output,
-                intermediate_tensors,
-                ubatch_slices,
-                num_tokens_after_padding,
-            )
-
-        cudagraph_runtime_mode = CUDAGraphMode.NONE
-        # Run the model.
-        # Use persistent buffers for CUDA graphs.
-        with (
-            set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_input_tokens,
-                num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_runtime_mode,
-                batch_descriptor=None,
-                ubatch_slices=ubatch_slices,
-            ),
-            record_function_or_nullcontext("Forward"),
-            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
-        ):
-
-            outputs = self._run_diffusion(
-                input_ids=input_ids,
-                positions=positions,
-                intermediate_tensors=intermediate_tensors,
-                inputs_embeds=inputs_embeds,
-                multimodal_kwargs=model_kwargs,
-                logits_indices=logits_indices,
-            )
-
-        _, multimodal_outputs = self.extract_multimodal_outputs(outputs)
-        # Ensure one tensor per request, map to CPU for output struct
-        pooler_output: List[Optional[torch.Tensor]] = []
-        if isinstance(multimodal_outputs, torch.Tensor):
-            # If model returned a single stacked tensor, split by requests
-            assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
-            for i in range(self.input_batch.num_reqs):
-                pooler_output.append(
-                    multimodal_outputs[i].detach().to("cpu").contiguous()
-                )
-        elif isinstance(multimodal_outputs, list):
-            for out in multimodal_outputs:
-                pooler_output.append(
-                    out.detach().to("cpu").contiguous() if out is not None else None
-                )
-        elif isinstance(multimodal_outputs, dict):
-            for out in multimodal_outputs.values():
-                pooler_output.append(
-                    out.detach().to("cpu").contiguous() if out is not None else None
-                )
-        else:
-            raise RuntimeError("Unsupported diffusion output type")
-
-        output = OmniModelRunnerOutput(
-            req_ids=self.input_batch.req_ids,
-            req_id_to_index=self.input_batch.req_id_to_index,
-            sampled_token_ids=[],
-            logprobs=None,
-            prompt_logprobs_dict={},
-            pooler_output=pooler_output,
-            kv_connector_output=kv_connector_output,
-            num_nans_in_logits={},
-        )
-
-        if not self.use_async_scheduling:
-            return output
-
-        return AsyncGPUModelRunnerOutput(
-            model_runner_output=output,
-            sampled_token_ids=[],
-            invalid_req_indices=[],
-            async_output_copy_stream=self.async_output_copy_stream,
-        )
-
-    def _run_diffusion(
-        self,
-        *,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        intermediate_tensors: Optional[IntermediateTensors],
-        inputs_embeds: Optional[torch.Tensor],
-        multimodal_kwargs: dict,
-        logits_indices: torch.Tensor,
-    ) -> Union[torch.Tensor, list[torch.Tensor]]:
-        """Runs the diffusion process and returns per-request tensors.
-
-        Tries model interfaces in the following order for maximal compatibility:
-        1) model.sample(condition=..., **kwargs)
-        2) model.forward(condition=..., **kwargs)
-        3) model.diffuse(condition=..., **kwargs)
-        """
-        # Keep inputs identical to AR runner
-        kwargs = dict(
-            input_ids=input_ids,
-            positions=positions,
-            intermediate_tensors=intermediate_tensors,
-            inputs_embeds=inputs_embeds,
-            **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
-            sampling_metadata=self.input_batch.sampling_metadata,
-            logits_index=logits_indices,
-            sampler=self.sampler,
-        )
-
-        # For Qwen 2.5 Omni's current implementation, we only support the forward method
-        if hasattr(self.model, "forward"):
-            return self.model.forward(**kwargs)
-        # TODO: add the diffuse method for other models
-
-        raise RuntimeError(
-            "The loaded model does not expose diffusion interfaces 'sample', "
-            "'forward', or 'diffuse'. Please implement one of them or adapt the runner."
-        )
-
-    @torch.inference_mode()
-    def _dummy_run(
-        self,
-        num_tokens: int,
-        cudagraph_runtime_mode: Optional[CUDAGraphMode] = None,
-        force_attention: bool = False,
-        uniform_decode: bool = False,
-        skip_eplb: bool = False,
-        is_profile: bool = False,
-        create_mixed_batch: bool = False,
-        remove_lora: bool = True,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Run a dummy forward pass to warm up/profile run or capture the
-        CUDA graph for the model.
-
-        Args:
-            num_tokens: Number of tokens to run the dummy forward pass.
-            cudagraph_runtime_mode: used to control the behavior.
-                - if not set will determine the cudagraph mode based on using
-                    the self.cudagraph_dispatcher.
-                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
-                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
-                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
-                    needed.
-            force_attention: If True, always create attention metadata. Used to
-                warm up attention backend when mode is NONE.
-            uniform_decode: If True, the batch is a uniform decode batch.
-            skip_eplb: If True, skip EPLB state update.
-            is_profile: If True, this is a profile run.
-            create_mixed_batch: If True, create a mixed batch with both decode
-                (1 token) and prefill (multiple tokens) requests.
-            remove_lora: If False, dummy LoRAs are not destroyed after the run
-        """
-        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
-            CUDAGraphMode.NONE,
-            CUDAGraphMode.PIECEWISE,
-            CUDAGraphMode.FULL,
-        }
-
-        # If cudagraph_mode.decode_mode() == FULL and
-        # cudagraph_mode.separate_routine(). This means that we are using
-        # different graphs and/or modes for mixed prefill-decode batches vs.
-        # uniform decode batches. A uniform decode batch means that all
-        # requests have identical query length, except a potential virtual
-        # request (shorter) in the batch account for padding.
-        # Uniform decode batch could either be common pure decode, where
-        # max_query_len == 1, or speculative decode, where
-        # max_query_len == 1 + num_spec_decode_tokens.
-
-        # When setting max_query_len = 1, we switch to and capture the optimized
-        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
-        # for GQA/MQA.
-        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
-
-        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
-        # for dummy run with LoRA so that the num_reqs collectively
-        # has num_tokens in total.
-        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
-        max_num_reqs = self.scheduler_config.max_num_seqs
-
-        num_reqs = min(num_tokens, max_num_reqs)
-        min_tokens_per_req = num_tokens // num_reqs
-        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
-        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
-
-        assert sum(num_scheduled_tokens_list) == num_tokens
-        assert len(num_scheduled_tokens_list) == num_reqs
-        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
-
-        num_tokens_after_padding = None
-
-        # If we failed to microbatch, currently need to resynchronize
-        # TODO(lucas,sage): we should be able to avoid this second sync by
-        #  refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
-        #  a single `coordinate_batch_across_dp` function.
-        if num_tokens_after_padding is None:
-            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
-            num_tokens_after_padding = num_tokens + num_pad
-        else:
-            num_tokens_across_dp = num_tokens_after_padding
-            num_tokens_after_padding = int(num_tokens_after_padding[0].item())
-
-        attn_metadata: Optional[PerLayerAttnMetadata] = None
-
-        # If force_attention is True, we always capture attention. Otherwise,
-        # it only happens for cudagraph_runtime_mode=FULL.
-        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
-            attn_metadata = {}
-
-            seq_lens = max_query_len
-            self.seq_lens.np[:num_reqs] = seq_lens
-            self.seq_lens.np[num_reqs:] = 0
-            self.seq_lens.copy_to_gpu()
-
-            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
-            self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
-            self.query_start_loc.copy_to_gpu()
-
-            for kv_cache_group_id, kv_cache_group_spec in enumerate(
-                self.kv_cache_config.kv_cache_groups
-            ):
-                common_attn_metadata = CommonAttentionMetadata(
-                    query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
-                    seq_lens=self.seq_lens.gpu[:num_reqs],
-                    seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
-                        :num_reqs
-                    ],
-                    num_reqs=num_reqs,
-                    num_actual_tokens=num_tokens,
-                    max_query_len=max_query_len,
-                    max_seq_len=self.max_model_len,
-                    block_table_tensor=self.input_batch.block_table[
-                        kv_cache_group_id
-                    ].get_device_tensor(num_reqs),
-                    slot_mapping=self.input_batch.block_table[
-                        kv_cache_group_id
-                    ].slot_mapping.gpu[:num_tokens],
-                    causal=True,
-                )
-                for attn_group in self.attn_groups[kv_cache_group_id]:
-
-                    assert type(attn_metadata) is dict
-                    attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(  # noqa: E501
-                        common_attn_metadata
-                    )
-                    for layer_name in attn_group.layer_names:
-                        attn_metadata[layer_name] = attn_metadata_i
-
-        with self.maybe_dummy_run_with_lora(
-            self.lora_config, num_scheduled_tokens, remove_lora
-        ):
-            model_kwargs = self._init_model_kwargs(num_tokens)
-            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
-                model_kwargs = {
-                    **model_kwargs,
-                    **self._dummy_mm_kwargs(num_reqs),
-                }
-            elif self.enable_prompt_embeds:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
-                model_kwargs = self._init_model_kwargs(num_tokens)
-            else:
-                input_ids = self.input_ids.gpu[:num_tokens]
-                inputs_embeds = None
-
-            if self.uses_mrope:
-                positions = self.mrope_positions.gpu[:, :num_tokens]
-            else:
-                positions = self.positions.gpu[:num_tokens]
-
-            if get_pp_group().is_first_rank:
-                intermediate_tensors = None
-            else:
-                if self.intermediate_tensors is None:
-                    self.intermediate_tensors = (
-                        self.model.make_empty_intermediate_tensors(
-                            batch_size=self.max_num_tokens,
-                            dtype=self.model_config.dtype,
-                            device=self.device,
-                        )
-                    )
-
-                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                    num_tokens, None, False
-                )
-
-            # filter out the valid batch descriptor
-            _cg_mode, batch_descriptor = (
-                self.cudagraph_dispatcher.dispatch(
-                    BatchDescriptor(
-                        num_tokens=num_tokens_after_padding,
-                        uniform_decode=uniform_decode,
-                    )
-                )
-                if not is_profile
-                else (CUDAGraphMode.NONE, None)
-            )
-            if cudagraph_runtime_mode is not None:
-                # we allow forcing NONE when the dispatcher disagrees to support
-                # warm ups for cudagraph capture
-                assert (
-                    cudagraph_runtime_mode == CUDAGraphMode.NONE
-                    or cudagraph_runtime_mode == _cg_mode
-                ), (
-                    f"Cudagraph runtime mode mismatch at dummy_run. "
-                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
-                )
-            else:
-                cudagraph_runtime_mode = _cg_mode
-
-            with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_tokens_after_padding,
-                num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_runtime_mode,
-                batch_descriptor=batch_descriptor,
-            ):
-                outputs = self.model(
-                    input_ids=input_ids,
-                    positions=positions,
-                    intermediate_tensors=intermediate_tensors,
-                    inputs_embeds=inputs_embeds,
-                    **model_kwargs,
-                )
-
-            if self.use_aux_hidden_state_outputs:
-                hidden_states, _ = outputs
-            else:
-                hidden_states = outputs
-
-            if self.speculative_config and self.speculative_config.use_eagle():
-                assert isinstance(self.drafter, EagleProposer)
-                self.drafter.dummy_run(num_tokens)
-
-        # This is necessary to avoid blocking DP.
-        # For dummy runs, we typically skip EPLB since we don't have any real
-        # requests to process.
-        # However, in DP settings, there may be cases when some DP ranks do
-        # not have any requests to process, so they're executing dummy batches.
-        # In such cases, we still have to trigger EPLB to make sure
-        # ranks execute the rearrangement in synchronization.
-
-        # logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
-        return hidden_states, None
-
-    @torch.inference_mode()
-    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
-        logger.warning("Dummy sampler run is not implemented for diffusion model")
-        return None
-
-    def profile_run(self) -> None:
-        # Profile with multimodal encoder & encoder cache.
-        if self.supports_mm_inputs:
-            if self.model_config.multimodal_config.skip_mm_profiling:
-                logger.info(
-                    "Skipping memory profiling for multimodal encoder and "
-                    "encoder cache."
-                )
-            else:
-                mm_budget = self.mm_budget
-                assert mm_budget is not None
-
-                if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
-                    # NOTE: Currently model is profiled with a single non-text
-                    # modality with the max possible input tokens even when
-                    # it supports multiple.
-                    dummy_modality = mm_budget.get_modality_with_max_tokens()
-                    max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[
-                        dummy_modality
-                    ]
-
-                    logger.info(
-                        "Encoder cache will be initialized with a budget of "
-                        "%s tokens, and profiled with %s %s items of the "
-                        "maximum feature size.",
-                        encoder_budget,
-                        max_mm_items_per_batch,
-                        dummy_modality,
-                    )
-
-                    # Create dummy batch of multimodal inputs.
-                    batched_dummy_mm_inputs = self._get_mm_dummy_batch(
-                        dummy_modality,
-                        max_mm_items_per_batch,
-                    )
-
-                    # Run multimodal encoder.
-                    dummy_encoder_outputs = self.model.get_multimodal_embeddings(
-                        **batched_dummy_mm_inputs
-                    )
-
-                    sanity_check_mm_encoder_outputs(
-                        dummy_encoder_outputs,
-                        expected_num_items=max_mm_items_per_batch,
-                    )
-
-                    # NOTE: This happens when encoder cache needs to store
-                    # the embeddings that encoder outputs are scattered onto.
-                    # In this case we create dummy embeddings of size
-                    # (encode_budget, hidden_size) and scatter encoder
-                    # output into it.
-                    encoder_output_shape = dummy_encoder_outputs[0].shape
-                    if encoder_output_shape[0] < encoder_budget:
-                        expanded_outputs = []
-                        for output in dummy_encoder_outputs:
-                            expanded = output.new_zeros(
-                                (encoder_budget, encoder_output_shape[-1])
-                            )
-                            num_tokens = output.shape[0]
-                            expanded[:num_tokens].copy_(output)
-                            expanded_outputs.append(expanded)
-
-                        dummy_encoder_outputs = expanded_outputs
-
-                    # Cache the dummy encoder outputs.
-                    self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
-
-        # Add `is_profile` here to pre-allocate communication buffers
-        hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
-        if get_pp_group().is_last_rank:
-            pass
-        self._sync_device()
-        del hidden_states
-        self.encoder_cache.clear()
-        gc.collect()

vllm_omni/worker/cpu/cpu_diffusion_worker.py

@@ -1,66 +0,0 @@
-import gc
-import os
-import platform
-
-import torch
-from vllm import envs
-from vllm.logger import init_logger
-
-from vllm.model_executor import set_random_seed
-from vllm.platforms import current_platform, CpuArchEnum
-
-from vllm.v1.worker.cpu_worker import CPUWorker
-
-from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-
-from vllm_omni.worker.cpu.cpu_diffusion_model_runner import CPUDiffusionModelRunner
-
-
-logger = init_logger(__name__)
-
-
-class CPUDiffusionWorker(CPUWorker):
-    def init_device(self):
-        # Setup OpenMP threads affinity.
-        omp_cpuids = envs.VLLM_CPU_OMP_THREADS_BIND
-        if omp_cpuids == "auto" and platform.system() == "Linux":
-            cpu_arch = current_platform.get_cpu_architecture()
-            if cpu_arch in (CpuArchEnum.POWERPC, CpuArchEnum.S390X):
-                # For S390X/POWERPC SMT-8/4/2
-                self.local_omp_cpuid = self._get_autobind_cpu_ids(
-                    lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4])
-            elif current_platform.get_cpu_architecture() == CpuArchEnum.X86:
-                # For x86 SMT-2, use 1 CPU per core
-                self.local_omp_cpuid = self._get_autobind_cpu_ids(
-                    lambda cpus: cpus[-1:])
-            else:
-                self.local_omp_cpuid = "all"
-        else:
-            local_dp_rank = self.parallel_config.data_parallel_rank_local
-            omp_cpuids = omp_cpuids.split("|")
-            if local_dp_rank is not None:
-                world_size = self.parallel_config.world_size
-                omp_cpuids = omp_cpuids[local_dp_rank *
-                                        world_size:(local_dp_rank + 1) *
-                                                   world_size]
-            self.local_omp_cpuid = omp_cpuids[self.rank]
-
-        if self.local_omp_cpuid != "all":
-            ret = torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)
-            if ret:
-                logger.info(ret)
-
-        # Note: unique identifier for creating allreduce shared memory
-        os.environ["VLLM_DIST_IDENT"] = self.distributed_init_method.split(
-            ":")[-1]
-        # Initialize the distributed environment.
-        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)
-
-        # Construct the model runner
-        self.model_runner: CPUDiffusionModelRunner = CPUDiffusionModelRunner(
-            self.vllm_config, torch.device("cpu"))

vllm_omni/worker/cpu/cpu_model_runner.py

@@ -1,593 +0,0 @@
-from typing import TYPE_CHECKING, List, Optional, Union, cast
-
-import numpy as np
-import torch
-
-from vllm.config import CUDAGraphMode
-from vllm.distributed.parallel_state import get_pp_group
-from vllm.forward_context import BatchDescriptor, set_forward_context
-from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces_base import VllmModelForPooling
-from vllm.sampling_params import SamplingType
-from vllm.utils import LazyLoader, cdiv
-from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata,
-    split_attn_metadata,
-)
-from vllm.v1.spec_decode.eagle import EagleProposer
-# from vllm.v1.worker.cpu import CachedRequestState
-from vllm.v1.worker.cpu_model_runner import CPUModelRunner
-from vllm.v1.worker.gpu_input_batch import CachedRequestState
-from vllm.v1.worker.gpu_model_runner import PerLayerAttnMetadata
-from vllm.v1.worker.ubatch_splitting import ubatch_split
-
-if TYPE_CHECKING:
-    from vllm.v1.core.sched.output import SchedulerOutput
-else:
-    xgr = LazyLoader("xgr", globals(), "xgrammar")
-    xgr_torch_compile = LazyLoader(
-        "xgr_torch_compile",
-        globals(),
-        "xgrammar.kernels.apply_token_bitmask_inplace_torch_compile",
-    )
-
-logger = init_logger(__name__)
-
-
-class OmniCPUModelRunner(CPUModelRunner):
-    def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
-        """Update the cached states and the persistent batch with the scheduler
-        output.
-
-        The updated states are used by the `_prepare_inputs` function to create
-        the input CPU tensors for the model.
-
-        The SamplingMetadata is updated and copied to the CPU if there is a
-        new/resumed/paused/finished request in the batch.
-        """
-        # Remove finished requests from the cached states.
-        for req_id in scheduler_output.finished_req_ids:
-            self.requests.pop(req_id, None)
-        # Remove the finished requests from the persistent batch.
-        # NOTE(woosuk): There could be an edge case where finished_req_ids and
-        # scheduled_req_ids overlap. This happens when a request is aborted and
-        # then resubmitted with the same ID. In this case, we treat them as two
-        # distinct requests - clearing the cached states for the first request
-        # and handling the second as a new request.
-        for req_id in scheduler_output.finished_req_ids:
-            self.input_batch.remove_request(req_id)
-
-        # Free the cached encoder outputs.
-        for mm_hash in scheduler_output.free_encoder_mm_hashes:
-            self.encoder_cache.pop(mm_hash, None)
-
-        # Remove the unscheduled requests from the persistent batch.
-        # NOTE(woosuk): The unscheduled requests are either preempted requests
-        # or running requests that are not scheduled in this step. We remove
-        # them from the persistent batch but keep their cached states since
-        # they will be scheduled again sometime in the future.
-        scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
-        cached_req_ids = self.input_batch.req_id_to_index.keys()
-        unscheduled_req_ids = cached_req_ids - scheduled_req_ids
-        # NOTE(woosuk): The persistent batch optimization assumes that
-        # consecutive batches contain mostly the same requests. If batches
-        # have low request overlap (e.g., alternating between two distinct
-        # sets of requests), this optimization becomes very inefficient.
-        for req_id in unscheduled_req_ids:
-            self.input_batch.remove_request(req_id)
-
-        reqs_to_add: list[CachedRequestState] = []
-        # Add new requests to the cached states.
-        for new_req_data in scheduler_output.scheduled_new_reqs:
-            req_id = new_req_data.req_id
-            sampling_params = new_req_data.sampling_params
-            pooling_params = new_req_data.pooling_params
-
-            if (
-                    sampling_params
-                    and sampling_params.sampling_type == SamplingType.RANDOM_SEED
-            ):
-                generator = torch.Generator(device=self.device)
-                generator.manual_seed(sampling_params.seed)
-            else:
-                generator = None
-
-            if self.is_pooling_model:
-                assert pooling_params is not None
-                task = pooling_params.task
-                assert task is not None, "You did not set `task` in the API"
-
-                model = cast(VllmModelForPooling, self.get_model())
-                to_update = model.pooler.get_pooling_updates(task)
-                to_update.apply(pooling_params)
-
-            req_state = CachedRequestState(
-                req_id=req_id,
-                prompt_token_ids=new_req_data.prompt_token_ids,
-                prompt_embeds=new_req_data.prompt_embeds,
-                mm_features=new_req_data.mm_features,
-                sampling_params=sampling_params,
-                pooling_params=pooling_params,
-                generator=generator,
-                block_ids=new_req_data.block_ids,
-                num_computed_tokens=new_req_data.num_computed_tokens,
-                output_token_ids=[],
-                lora_request=new_req_data.lora_request,
-            )
-            self.requests[req_id] = req_state
-
-            # If prompt embeddings are provided, decode and attach to inter_data
-            try:
-                if getattr(new_req_data, "prompt_embeds", None) is not None:
-                    payload = new_req_data.prompt_embeds
-                    dtype = getattr(np, payload.dtype)
-                    arr = np.frombuffer(payload.data, dtype=dtype)
-                    arr = arr.reshape(payload.shape)
-                    pe_cpu = torch.from_numpy(arr)
-                    # Store temporarily on CPU; later moved to device in builder
-                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
-                    # Also replace payload with Tensor for user visibility in
-                    # scheduler_output
-                    try:
-                        new_req_data.prompt_embeds = pe_cpu  # type: ignore[assignment]
-                    except Exception:
-                        pass
-            except Exception as e:
-                logger.error(f"Error decoding prompt embeds: {e}")
-            # Decode additional_information payloads (dictionary)
-            try:
-                if getattr(new_req_data, "additional_information", None) is not None:
-                    payload_info = new_req_data.additional_information
-                    info_dict = {}
-                    if isinstance(payload_info, dict):
-                        info_dict = payload_info
-                    else:
-                        from vllm_omni.engine import AdditionalInformationPayload
-
-                        if isinstance(payload_info, AdditionalInformationPayload):
-                            for k, entry in payload_info.entries.items():
-                                if entry.tensor_data is not None:
-                                    dt = np.dtype(
-                                        getattr(entry, "tensor_dtype", "float32")
-                                    )
-                                    arr = np.frombuffer(entry.tensor_data, dtype=dt)
-                                    arr = arr.reshape(entry.tensor_shape)
-                                    info_dict[k] = torch.from_numpy(arr)
-                                else:
-                                    info_dict[k] = entry.list_data
-                    if info_dict:
-                        setattr(
-                            self.requests[req_id],
-                            "additional_information_cpu",
-                            info_dict,
-                        )
-            except Exception as e:
-                logger.error(f"Error decoding additional information: {e}")
-                pass
-
-            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
-            if self.uses_mrope:
-                self._init_mrope_positions(req_state)
-
-            reqs_to_add.append(self.requests[req_id])
-
-        # Update the states of the running/resumed requests.
-        is_last_rank = get_pp_group().is_last_rank
-        req_data = scheduler_output.scheduled_cached_reqs
-        for i, req_id in enumerate(req_data.req_ids):
-            req_state = self.requests[req_id]
-            num_computed_tokens = req_data.num_computed_tokens[i]
-            new_block_ids = req_data.new_block_ids[i]
-            resumed_from_preemption = req_data.resumed_from_preemption[i]
-
-            # Update the cached states.
-            req_state.num_computed_tokens = num_computed_tokens
-
-            if not is_last_rank:
-                # When using PP, the scheduler sends the sampled tokens back,
-                # because there's no direct communication between the first-
-                # stage worker and the last-stage worker.
-                new_token_ids = req_data.new_token_ids[i]
-                # Add the sampled token(s) from the previous step (if any).
-                # This doesn't include "unverified" tokens like spec tokens.
-                num_new_tokens = (
-                        num_computed_tokens + len(new_token_ids) - req_state.num_tokens
-                )
-                if num_new_tokens == 1:
-                    # Avoid slicing list in most common case.
-                    req_state.output_token_ids.append(new_token_ids[-1])
-                elif num_new_tokens > 0:
-                    req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
-
-            # Update the block IDs.
-            if not resumed_from_preemption:
-                if new_block_ids is not None:
-                    # Append the new blocks to the existing block IDs.
-                    for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
-                        block_ids.extend(new_ids)
-            else:
-                assert new_block_ids is not None
-                # The request is resumed from preemption.
-                # Replace the existing block IDs with the new ones.
-                req_state.block_ids = new_block_ids
-
-            req_index = self.input_batch.req_id_to_index.get(req_id)
-            if req_index is None:
-                # The request is not in the persistent batch.
-                # The request was either preempted and resumed later, or was not
-                # scheduled in the previous step and needs to be added again.
-                reqs_to_add.append(req_state)
-                continue
-
-            # Update the persistent batch.
-            self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens
-            if new_block_ids is not None:
-                self.input_batch.block_table.append_row(new_block_ids, req_index)
-
-            # For the last rank, we don't need to update the token_ids_cpu
-            # because the sampled tokens are already cached.
-            if not is_last_rank:
-                # Add new_token_ids to token_ids_cpu.
-                start_token_index = num_computed_tokens
-                end_token_index = num_computed_tokens + len(new_token_ids)
-                self.input_batch.token_ids_cpu[
-                    req_index, start_token_index:end_token_index
-                ] = new_token_ids
-                self.input_batch.num_tokens_no_spec[req_index] = end_token_index
-                self.input_batch.num_tokens[req_index] = end_token_index
-
-            # Add spec_token_ids to token_ids_cpu.
-            spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
-                req_id, ()
-            )
-            if spec_token_ids:
-                num_spec_tokens = len(spec_token_ids)
-                start_index = self.input_batch.num_tokens_no_spec[req_index]
-                end_token_index = start_index + num_spec_tokens
-                self.input_batch.token_ids_cpu[
-                    req_index, start_index:end_token_index
-                ] = spec_token_ids
-                # NOTE(woosuk): `num_tokens` here may include spec tokens.
-                self.input_batch.num_tokens[req_index] += num_spec_tokens
-
-        # Add the new or resumed requests to the persistent batch.
-        # The smaller empty indices are filled first.
-        for request in reqs_to_add:
-            self.input_batch.add_request(request)
-
-        # Condense the batched states if there are gaps left by removed requests
-        self.input_batch.condense()
-        # Allow attention backend to reorder the batch, potentially
-        self._may_reorder_batch(scheduler_output)
-        # Refresh batch metadata with any pending updates.
-        self.input_batch.refresh_metadata()
-
-    @torch.inference_mode()
-    def extract_multimodal_outputs(
-            self, hidden_states: Union[torch.Tensor, List[torch.Tensor]]
-    ) -> dict:
-        if (
-                hasattr(self.model, "have_multimodal_outputs")
-                and self.model.have_multimodal_outputs
-        ):
-            text_hidden_states = hidden_states.text_hidden_states
-            multimodal_outputs = hidden_states.multimodal_outputs
-
-        elif isinstance(hidden_states, torch.Tensor):
-            text_hidden_states = hidden_states
-            multimodal_outputs = {}
-        elif isinstance(hidden_states, List):
-            text_hidden_states = hidden_states[0]
-            multimodal_outputs = {}
-        else:
-            raise ValueError(f"Invalid hidden states type: {type(hidden_states)}")
-        return text_hidden_states, multimodal_outputs
-
-    @torch.inference_mode()
-    def _dummy_run(
-            self,
-            num_tokens: int,
-            cudagraph_runtime_mode: Optional[CUDAGraphMode] = None,
-            force_attention: bool = False,
-            uniform_decode: bool = False,
-            allow_microbatching: bool = True,
-            skip_eplb: bool = False,
-            is_profile: bool = False,
-            create_mixed_batch: bool = False,
-            remove_lora: bool = True,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Run a dummy forward pass to warm up/profile run or capture the
-        CUDA graph for the model.
-
-        Args:
-            num_tokens: Number of tokens to run the dummy forward pass.
-            cudagraph_runtime_mode: used to control the behavior.
-                - if not set will determine the cudagraph mode based on using
-                    the self.cudagraph_dispatcher.
-                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
-                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
-                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
-                    needed.
-            force_attention: If True, always create attention metadata. Used to
-                warm up attention backend when mode is NONE.
-            uniform_decode: If True, the batch is a uniform decode batch.
-            skip_eplb: If True, skip EPLB state update.
-            is_profile: If True, this is a profile run.
-            create_mixed_batch: If True, create a mixed batch with both decode
-                (1 token) and prefill (multiple tokens) requests.
-            remove_lora: If False, dummy LoRAs are not destroyed after the run
-        """
-        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
-            CUDAGraphMode.NONE,
-            CUDAGraphMode.PIECEWISE,
-            CUDAGraphMode.FULL,
-        }
-
-        # If cudagraph_mode.decode_mode() == FULL and
-        # cudagraph_mode.separate_routine(). This means that we are using
-        # different graphs and/or modes for mixed prefill-decode batches vs.
-        # uniform decode batches. A uniform decode batch means that all
-        # requests have identical query length, except a potential virtual
-        # request (shorter) in the batch account for padding.
-        # Uniform decode batch could either be common pure decode, where
-        # max_query_len == 1, or speculative decode, where
-        # max_query_len == 1 + num_spec_decode_tokens.
-
-        # When setting max_query_len = 1, we switch to and capture the optimized
-        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
-        # for GQA/MQA.
-        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
-
-        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
-        # for dummy run with LoRA so that the num_reqs collectively
-        # has num_tokens in total.
-        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
-        max_num_reqs = self.scheduler_config.max_num_seqs
-        if create_mixed_batch:
-            assert not uniform_decode
-            # Create mixed batch:
-            # first half decode tokens, second half one prefill
-            num_decode_tokens = num_tokens // 2
-            num_prefill_tokens = num_tokens - num_decode_tokens
-            num_reqs = num_decode_tokens + 1
-
-            # Create decode requests (1 token each) followed by prefill request
-            num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens]
-            # Note: Overriding max_query_len to be the prefill tokens
-            max_query_len = num_prefill_tokens
-        elif uniform_decode:
-            assert not create_mixed_batch
-            num_reqs = cdiv(num_tokens, max_query_len)
-            num_scheduled_tokens_list = [max_query_len] * num_reqs
-            if num_tokens % max_query_len != 0:
-                num_scheduled_tokens_list[-1] = num_tokens % max_query_len
-        else:
-            num_reqs = min(num_tokens, max_num_reqs)
-            min_tokens_per_req = num_tokens // num_reqs
-            num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
-            num_scheduled_tokens_list[-1] += num_tokens % num_reqs
-
-        assert sum(num_scheduled_tokens_list) == num_tokens
-        assert len(num_scheduled_tokens_list) == num_reqs
-        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
-
-        total_num_scheduled_tokens = int(num_scheduled_tokens.sum())
-
-        ubatch_slices = None
-        num_tokens_after_padding = None
-
-        # We currently only microbatch if the number of tokens is
-        # over a certain threshold.
-        if self.parallel_config.enable_dbo and allow_microbatching:
-            ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split(
-                num_scheduled_tokens,
-                total_num_scheduled_tokens,
-                total_num_scheduled_tokens,
-                uniform_decode=uniform_decode,
-                vllm_config=self.vllm_config,
-            )
-            # Currently when DBO is enabled `ubatch_split` returns
-            # the num_tokens_after_padding for a single ubatch, but we have 2
-            # TODO(sage,lucas): this is cruft that should be addressed in the
-            # padding refactor.
-            if ubatch_num_tokens_after_padding is not None:
-                num_tokens_after_padding = ubatch_num_tokens_after_padding * 2
-
-        # If we failed to microbatch, currently need to resynchronize
-        # TODO(lucas,sage): we should be able to avoid this second sync by
-        #  refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
-        #  a single `coordinate_batch_across_dp` function.
-        if num_tokens_after_padding is None:
-            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
-            num_tokens_after_padding = num_tokens + num_pad
-        else:
-            num_tokens_across_dp = num_tokens_after_padding
-            num_tokens_after_padding = int(num_tokens_after_padding[0].item())
-
-        attn_metadata: Optional[PerLayerAttnMetadata] = None
-
-        # If force_attention is True, we always capture attention. Otherwise,
-        # it only happens for cudagraph_runtime_mode=FULL.
-        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
-            attn_metadata = {}
-            if ubatch_slices is not None:
-                attn_metadata = [dict() for _ in range(len(ubatch_slices))]
-
-            if create_mixed_batch:
-                # In the mixed batch mode (used for FI warmup), we use
-                # shorter sequence lengths to run faster.
-                # TODO(luka) better system for describing dummy batches
-                seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
-            else:
-                seq_lens = max_query_len
-            self.seq_lens.np[:num_reqs] = seq_lens
-            self.seq_lens.np[num_reqs:] = 0
-            self.seq_lens.copy_to_gpu()
-
-            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
-            self.query_start_loc.np[1: num_reqs + 1] = cum_num_tokens
-            self.query_start_loc.copy_to_gpu()
-
-            for kv_cache_group_id, kv_cache_group_spec in enumerate(
-                    self.kv_cache_config.kv_cache_groups
-            ):
-                common_attn_metadata = CommonAttentionMetadata(
-                    query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
-                    seq_lens=self.seq_lens.gpu[:num_reqs],
-                    seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
-                        :num_reqs
-                    ],
-                    num_reqs=num_reqs,
-                    num_actual_tokens=num_tokens,
-                    max_query_len=max_query_len,
-                    max_seq_len=self.max_model_len,
-                    block_table_tensor=self.input_batch.block_table[
-                        kv_cache_group_id
-                    ].get_device_tensor(num_reqs),
-                    slot_mapping=self.input_batch.block_table[
-                        kv_cache_group_id
-                    ].slot_mapping.gpu[:num_tokens],
-                    causal=True,
-                )
-                for attn_group in self.attn_groups[kv_cache_group_id]:
-                    if ubatch_slices is not None:
-                        common_attn_metadata_list = split_attn_metadata(
-                            ubatch_slices, common_attn_metadata
-                        )
-                        for ubid, common_attn_metadata in enumerate(
-                                common_attn_metadata_list
-                        ):
-                            assert common_attn_metadata.max_query_len == 1
-                            attn_metadata_i = attn_group.get_metadata_builder(
-                                ubatch_id=ubid
-                            ).build_for_cudagraph_capture(common_attn_metadata)
-                            for layer_name in attn_group.layer_names:
-                                assert type(attn_metadata) is list
-                                attn_metadata[ubid][layer_name] = attn_metadata_i
-                    else:
-                        assert type(attn_metadata) is dict
-                        attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(  # noqa: E501
-                            common_attn_metadata
-                        )
-                        for layer_name in attn_group.layer_names:
-                            attn_metadata[layer_name] = attn_metadata_i
-
-        with self.maybe_dummy_run_with_lora(
-                self.lora_config, num_scheduled_tokens, remove_lora
-        ):
-            model_kwargs = self._init_model_kwargs(num_tokens)
-            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
-                model_kwargs = {
-                    **model_kwargs,
-                    **self._dummy_mm_kwargs(num_reqs),
-                }
-            elif self.enable_prompt_embeds:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
-                model_kwargs = self._init_model_kwargs(num_tokens)
-            else:
-                input_ids = self.input_ids.gpu[:num_tokens]
-                inputs_embeds = None
-
-            if self.uses_mrope:
-                positions = self.mrope_positions.gpu[:, :num_tokens]
-            else:
-                positions = self.positions.gpu[:num_tokens]
-
-            if get_pp_group().is_first_rank:
-                intermediate_tensors = None
-            else:
-                if self.intermediate_tensors is None:
-                    self.intermediate_tensors = (
-                        self.model.make_empty_intermediate_tensors(
-                            batch_size=self.max_num_tokens,
-                            dtype=self.model_config.dtype,
-                            device=self.device,
-                        )
-                    )
-
-                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                    num_tokens, None, False
-                )
-
-            # filter out the valid batch descriptor
-            _cg_mode, batch_descriptor = (
-                self.cudagraph_dispatcher.dispatch(
-                    BatchDescriptor(
-                        num_tokens=num_tokens_after_padding,
-                        uniform_decode=uniform_decode,
-                    )
-                )
-                if not is_profile
-                else (CUDAGraphMode.NONE, None)
-            )
-            if cudagraph_runtime_mode is not None:
-                # we allow forcing NONE when the dispatcher disagrees to support
-                # warm ups for cudagraph capture
-                assert (
-                        cudagraph_runtime_mode == CUDAGraphMode.NONE
-                        or cudagraph_runtime_mode == _cg_mode
-                ), (
-                    f"Cudagraph runtime mode mismatch at dummy_run. "
-                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
-                )
-            else:
-                cudagraph_runtime_mode = _cg_mode
-
-            if ubatch_slices is not None:
-                # Adjust values to reflect a single ubatch.
-                # TODO(sage,lucas): this is cruft that should be addressed in
-                #  the padding refactor.
-                num_tokens_after_padding = ubatch_slices[0].num_tokens
-                if num_tokens_across_dp is not None:
-                    num_tokens_across_dp[:] = num_tokens_after_padding
-
-            with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                    attn_metadata,
-                    self.vllm_config,
-                    num_tokens=num_tokens_after_padding,
-                    num_tokens_across_dp=num_tokens_across_dp,
-                    cudagraph_runtime_mode=cudagraph_runtime_mode,
-                    batch_descriptor=batch_descriptor,
-                    ubatch_slices=ubatch_slices,
-            ):
-                outputs = self.model(
-                    input_ids=input_ids,
-                    positions=positions,
-                    intermediate_tensors=intermediate_tensors,
-                    inputs_embeds=inputs_embeds,
-                    **model_kwargs,
-                )
-
-            if self.use_aux_hidden_state_outputs:
-                hidden_states, _ = outputs
-            else:
-                hidden_states = outputs
-
-            if self.speculative_config and self.speculative_config.use_eagle():
-                assert isinstance(self.drafter, EagleProposer)
-                self.drafter.dummy_run(num_tokens)
-
-        # This is necessary to avoid blocking DP.
-        # For dummy runs, we typically skip EPLB since we don't have any real
-        # requests to process.
-        # However, in DP settings, there may be cases when some DP ranks do
-        # not have any requests to process, so they're executing dummy batches.
-        # In such cases, we still have to trigger EPLB to make sure
-        # ranks execute the rearrangement in synchronization.
-        if not skip_eplb:
-            self.eplb_step(is_dummy=True, is_profile=is_profile)
-
-        logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
-            hidden_states
-        )
-        if len(hidden_states.shape) == 3:
-            return hidden_states[0], hidden_states[0][logit_indices]
-        else:
-            return hidden_states, hidden_states[logit_indices]

vllm_omni/worker/gpu_ar_model_runner.py

@@ -6,11 +6,10 @@ and also outputs sampled tokens.
 
 from __future__ import annotations
 
-from typing import Any, Optional, Union
+from typing import Any
 
 import numpy as np
 import torch
-
 from vllm.forward_context import BatchDescriptor
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
@@ -28,9 +27,10 @@ from vllm.v1.worker.gpu_model_runner import (
 )
 from vllm.v1.worker.ubatch_utils import UBatchSlices
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
+
 from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.gpu.gpu_model_runner import OmniGPUModelRunner
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
 logger = init_logger(__name__)
 
@@ -51,23 +51,22 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
     def _preprocess(
         self,
-        scheduler_output: "SchedulerOutput",
+        scheduler_output: SchedulerOutput,
         num_scheduled_tokens_np: np.ndarray,
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-        ubatch_slices: Optional[UBatchSlices] = None,
-        num_tokens_after_padding: Optional[torch.Tensor] = None,
+        intermediate_tensors: IntermediateTensors | None = None,
+        ubatch_slices: UBatchSlices | None = None,
+        num_tokens_after_padding: torch.Tensor | None = None,
     ) -> tuple[
         int,
         int,
-        Optional[torch.Tensor],
-        Optional[torch.Tensor],
-        Optional[torch.Tensor],
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
         torch.Tensor,
-        Optional[IntermediateTensors],
+        IntermediateTensors | None,
         dict[str, Any],
-        Optional[dict[str, dict]],
+        dict[str, dict] | None,
     ]:
-
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         if ubatch_slices:
             assert num_tokens_after_padding is not None
@@ -80,12 +79,8 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
         # _prepare_inputs may reorder the batch, so we must gather multi
         # modal outputs after that to ensure the correct order
-        per_req_additional_information: Optional[dict[str, dict]] = None
-        if (
-            self.supports_mm_inputs
-            and get_pp_group().is_first_rank
-            and not self.model_config.is_encoder_decoder
-        ):
+        per_req_additional_information: dict[str, dict] | None = None
+        if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
             # Run the multimodal encoder if any.
             self._execute_mm_encoder(scheduler_output)
             mm_embeds = self._gather_mm_embeddings(scheduler_output)
@@ -113,9 +108,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 req_state = self.requests[req_id]
                 pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
                 addi_cpu = getattr(req_state, "additional_information_cpu", None)
-                num_computed_tokens = int(
-                    self.input_batch.num_computed_tokens_cpu[req_index]
-                )
+                num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index])
                 prompt_len = len(req_state.prompt_token_ids)
                 prompt_remaining = max(0, prompt_len - num_computed_tokens)
                 sched_tokens = int(num_scheduled_tokens_np[req_index])
@@ -123,24 +116,26 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 if overlay_len <= 0:
                     continue
                 if pe_cpu is not None:
-                    src = pe_cpu[
-                        num_computed_tokens : num_computed_tokens + overlay_len
-                    ].to(dtype=self.dtype, device=self.device, non_blocking=True)
-                    start_offset = int(self.query_start_loc.cpu[req_index])
-                    self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(
-                        src
+                    src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to(
+                        dtype=self.dtype, device=self.device, non_blocking=True
                     )
+                    start_offset = int(self.query_start_loc.cpu[req_index])
+                    self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
                 # Build per-request additional information (no cross-request concat)
                 if addi_cpu is not None and isinstance(addi_cpu, dict):
                     req_info: dict[str, object] = {}
                     for k, v in addi_cpu.items():
                         if isinstance(v, torch.Tensor):
-                            # For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor
+                            # For prefill tokens, pass only the scheduled slice;
+                            # for decode or no scheduled tokens, pass whole tensor
                             if overlay_len > 0:
                                 try:
-                                    seg = v[
-                                        num_computed_tokens : num_computed_tokens + overlay_len
-                                    ].detach().to("cpu").contiguous()
+                                    seg = (
+                                        v[num_computed_tokens : num_computed_tokens + overlay_len]
+                                        .detach()
+                                        .to("cpu")
+                                        .contiguous()
+                                    )
                                 except Exception:
                                     seg = v.detach().to("cpu").contiguous()
                                 req_info[k] = seg
@@ -171,11 +166,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             # If a batch only has token ids, then including the embedding layer
             # in the CUDA graph will be more performant (like in the else case
             # below).
-            token_ids_idx = (
-                self.is_token_ids.gpu[:num_scheduled_tokens]
-                .nonzero(as_tuple=False)
-                .squeeze(1)
-            )
+            token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1)
             # Some tokens ids may need to become embeds
             if token_ids_idx.numel() > 0:
                 token_ids = self.input_ids.gpu[token_ids_idx]
@@ -205,10 +196,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 num_input_tokens, intermediate_tensors, True
             )
 
-        if (
-            self.model_config.is_encoder_decoder
-            and scheduler_output.scheduled_encoder_inputs
-        ):
+        if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
             encoder_inputs = self._extract_encoder_inputs(scheduler_output)
             model_kwargs.update(encoder_inputs)
 
@@ -227,9 +215,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
     @torch.inference_mode()
     def execute_model(
         self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-    ) -> Union[OmniModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+    ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
         with record_function_or_nullcontext("Preprocess"):
             with self.synchronize_input_prep():
                 # Update persistent batch states.
@@ -241,9 +229,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                     new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
                     if new_reqs:
                         for nr in new_reqs:
-                            req_id = getattr(nr, "req_id", None) or getattr(
-                                nr, "request_id", None
-                            )
+                            req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
                             if req_id is None:
                                 continue
                             # prompt_embeds
@@ -252,9 +238,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                                 if isinstance(payload_pe, torch.Tensor):
                                     pe_cpu = payload_pe.detach().to("cpu").contiguous()
                                 elif isinstance(payload_pe, PromptEmbedsPayload):
-                                    dt = np.dtype(
-                                        getattr(payload_pe, "dtype", "float32")
-                                    )
+                                    dt = np.dtype(getattr(payload_pe, "dtype", "float32"))
                                     arr = np.frombuffer(payload_pe.data, dtype=dt)
                                     arr = arr.reshape(payload_pe.shape)
                                     pe_cpu = torch.from_numpy(arr)
@@ -273,19 +257,11 @@ class GPUARModelRunner(OmniGPUModelRunner):
                                 if isinstance(payload_info, dict):
                                     # Already decoded
                                     info_dict = payload_info
-                                elif isinstance(
-                                    payload_info, AdditionalInformationPayload
-                                ):
+                                elif isinstance(payload_info, AdditionalInformationPayload):
                                     for k, entry in payload_info.entries.items():
                                         if entry.tensor_data is not None:
-                                            dt = np.dtype(
-                                                getattr(
-                                                    entry, "tensor_dtype", "float32"
-                                                )
-                                            )
-                                            arr = np.frombuffer(
-                                                entry.tensor_data, dtype=dt
-                                            )
+                                            dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
+                                            arr = np.frombuffer(entry.tensor_data, dtype=dt)
                                             arr = arr.reshape(entry.tensor_shape)
                                             info_dict[k] = torch.from_numpy(arr)
                                         else:
@@ -297,18 +273,14 @@ class GPUARModelRunner(OmniGPUModelRunner):
                                         info_dict,
                                     )
                 except Exception as e:
-                    logger.error(
-                        f"Error decoding prompt_embeds / additional_information: {e}"
-                    )
+                    logger.error(f"Error decoding prompt_embeds / additional_information: {e}")
                     pass
 
                 if not scheduler_output.total_num_scheduled_tokens:
                     if not has_kv_transfer_group():
                         # Return empty ModelRunnerOutput if no work to do.
                         return EMPTY_MODEL_RUNNER_OUTPUT
-                    return self.kv_connector_no_forward(
-                        scheduler_output, self.vllm_config
-                    )
+                    return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
                 if self.cache_config.kv_sharing_fast_prefill:
                     assert not self.input_batch.num_prompt_logprobs, (
                         "--kv-sharing-fast-prefill produces incorrect "
@@ -349,12 +321,8 @@ class GPUARModelRunner(OmniGPUModelRunner):
             uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
                 num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
             )
-            batch_descriptor = BatchDescriptor(
-                num_tokens=num_input_tokens, uniform_decode=uniform_decode
-            )
-            cudagraph_runtime_mode, batch_descriptor = (
-                self.cudagraph_dispatcher.dispatch(batch_descriptor)
-            )
+            batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens, uniform_decode=uniform_decode)
+            cudagraph_runtime_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(batch_descriptor)
 
         # This is currently to get around the assert in the DPMetadata
         # where it wants `num_tokens_across_dp` to align with `num_tokens`
@@ -376,7 +344,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
             record_function_or_nullcontext("Forward"),
             self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
         ):
-
             model_kwargs_extra = {}
             # Pass per-request additional information map for this step (no concat)
             if per_req_additional_information:
@@ -386,17 +353,12 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 per_req_runtime_info = []
                 for req_id in self.input_batch.req_ids:
                     req_state = self.requests.get(req_id)
-                    info = (
-                        getattr(req_state, "additional_information_cpu", None)
-                        if req_state is not None
-                        else None
-                    )
+                    info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
                     per_req_runtime_info.append(info if isinstance(info, dict) else {})
-                model_kwargs_extra["runtime_additional_information"] = (
-                    per_req_runtime_info
-                )
+                model_kwargs_extra["runtime_additional_information"] = per_req_runtime_info
                 model_kwargs_extra["request_ids"] = self.input_batch.req_ids
-                # Pass each request's token span within the flattened sequence for this step, enabling the model to map decode/prefill by request
+                # Pass each request's token span within the flattened sequence for this step,
+                # enabling the model to map decode/prefill by request
                 req_token_spans = []
                 for req_index in range(len(self.input_batch.req_ids)):
                     start_offset = int(self.query_start_loc.cpu[req_index])
@@ -426,24 +388,17 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 hidden_states = model_output
                 aux_hidden_states = None
 
-            hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
-                hidden_states
-            )
+            hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
             # The model side may return per-request additional_information updates (model-agnostic channel).
             # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
             # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
             try:
-                if (
-                    isinstance(multimodal_outputs, dict)
-                    and (
-                        "additional_information_update" in multimodal_outputs
-                        or "additional_information_update_by_req_id" in multimodal_outputs
-                    )
+                if isinstance(multimodal_outputs, dict) and (
+                    "additional_information_update" in multimodal_outputs
+                    or "additional_information_update_by_req_id" in multimodal_outputs
                 ):
                     # Option A: list[dict] in batch order
-                    updates_list = multimodal_outputs.get(
-                        "additional_information_update"
-                    )
+                    updates_list = multimodal_outputs.get("additional_information_update")
                     if isinstance(updates_list, list):
                         for idx, upd in enumerate(updates_list):
                             if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
@@ -451,9 +406,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                             req_id = self.input_batch.req_ids[idx]
                             self._merge_additional_information_update(req_id, upd)
                     # Option B: dict[str, dict] keyed by req_id
-                    updates_map = multimodal_outputs.get(
-                        "additional_information_update_by_req_id"
-                    )
+                    updates_map = multimodal_outputs.get("additional_information_update_by_req_id")
                     if isinstance(updates_map, dict):
                         for req_id, upd in updates_map.items():
                             if not isinstance(upd, dict):
@@ -462,9 +415,11 @@ class GPUARModelRunner(OmniGPUModelRunner):
                                 continue
                             self._merge_additional_information_update(req_id, upd)
             except Exception as e:
-                logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}")
+                logger.error(
+                    f"Error merging for requests:{self.input_batch.req_ids} additional \
+                        information update: {e}, with the multimodal_outputs as {multimodal_outputs}"
+                )
             if not self.broadcast_pp_output:
-
                 # Common case.
                 if not get_pp_group().is_last_rank:
                     # Return the intermediate tensors.
@@ -474,9 +429,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
                 if self.is_pooling_model:
                     # Return the pooling output.
-                    output = self._pool(
-                        hidden_states, num_scheduled_tokens, num_scheduled_tokens_np
-                    )
+                    output = self._pool(hidden_states, num_scheduled_tokens, num_scheduled_tokens_np)
                     output.kv_connector_output = kv_connector_output
                     return output
 
@@ -488,9 +441,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
                 if not get_pp_group().is_last_rank:
                     all_gather_tensors = {
-                        "residual": not is_residual_scattered_for_sp(
-                            self.vllm_config, num_input_tokens
-                        )
+                        "residual": not is_residual_scattered_for_sp(self.vllm_config, num_input_tokens)
                     }
                     get_pp_group().send_tensor_dict(
                         hidden_states.tensors,
@@ -514,9 +465,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
             # Apply structured output bitmasks if present
             if scheduler_output.grammar_bitmask is not None:
-                apply_grammar_bitmask(
-                    scheduler_output, self.input_batch, logits, self.device
-                )
+                apply_grammar_bitmask(scheduler_output, self.input_batch, logits, self.device)
 
         with record_function_or_nullcontext("Sample"):
             sampler_output = self._sample(logits, spec_decode_metadata)
@@ -548,12 +497,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
             and self.speculative_config.draft_model_config is not None
             and self.speculative_config.draft_model_config.max_model_len is not None
         ):
-            effective_drafter_max_model_len = (
-                self.speculative_config.draft_model_config.max_model_len
-            )
+            effective_drafter_max_model_len = self.speculative_config.draft_model_config.max_model_len
         input_fits_in_drafter = spec_decode_common_attn_metadata and (
-            spec_decode_common_attn_metadata.seq_lens.max()
-            + self.speculative_config.num_speculative_tokens
+            spec_decode_common_attn_metadata.seq_lens.max() + self.speculative_config.num_speculative_tokens
             <= effective_drafter_max_model_len
         )
         if use_padded_batch_for_eagle and input_fits_in_drafter:
@@ -578,11 +524,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 num_scheduled_tokens,
             )
 
-        if (
-            self.speculative_config
-            and not use_padded_batch_for_eagle
-            and input_fits_in_drafter
-        ):
+        if self.speculative_config and not use_padded_batch_for_eagle and input_fits_in_drafter:
             # ngram and other speculative decoding methods use the sampled
             # tokens on the CPU, so they are run after bookkeeping.
             propose_draft_token_ids(valid_sampled_token_ids)
@@ -592,12 +534,10 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
         # Convert to per-request tensors on CPU
         hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
-        pooler_output: list[Optional[torch.Tensor]] = []
+        pooler_output: list[torch.Tensor | None] = []
         prev_logits_index = 0
         for logits_index in logits_indices:
-            pooler_output.append(
-                hidden_states_cpu[prev_logits_index : logits_index + 1]
-            )
+            pooler_output.append(hidden_states_cpu[prev_logits_index : logits_index + 1])
             prev_logits_index = logits_index + 1
 
         output = OmniModelRunnerOutput(
@@ -606,11 +546,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             sampled_token_ids=valid_sampled_token_ids,
             logprobs=logprobs_lists,
             prompt_logprobs_dict=prompt_logprobs_dict,
-            pooler_output=(
-                pooler_output
-                if self.vllm_config.model_config.engine_output_type != "text"
-                else None
-            ),
+            pooler_output=(pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None),
             kv_connector_output=kv_connector_output,
             num_nans_in_logits=num_nans_in_logits,
         )

vllm_omni/worker/gpu_ar_worker.py

@@ -2,14 +2,14 @@ import gc
 import os
 
 import torch
-
 from vllm.model_executor import set_random_seed
 from vllm.platforms import current_platform
 from vllm.utils import GiB_bytes, MemorySnapshot
 from vllm.v1.utils import report_usage_stats
 from vllm.v1.worker.gpu_worker import Worker as GPUWorker
 from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-from vllm_omni.worker.gpu.gpu_ar_model_runner import GPUARModelRunner
+
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
 
 
 class GPUARWorker(GPUWorker):
@@ -43,10 +43,7 @@ class GPUARWorker(GPUWorker):
 
             # take current memory snapshot
             self.init_snapshot = MemorySnapshot()
-            self.requested_memory = (
-                self.init_snapshot.total_memory
-                * self.cache_config.gpu_memory_utilization
-            )
+            self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization
             if self.init_snapshot.free_memory < self.requested_memory:
 
                 def GiB(b):
@@ -65,9 +62,7 @@ class GPUARWorker(GPUWorker):
             raise RuntimeError(f"Not support device type: {self.device_config.device}")
 
         # Construct the model runner
-        self.model_runner: GPUARModelRunner = GPUARModelRunner(
-            self.vllm_config, self.device
-        )
+        self.model_runner: GPUARModelRunner = GPUARModelRunner(self.vllm_config, self.device)
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.

vllm_omni/worker/gpu_diffusion_model_runner.py

@@ -2,11 +2,10 @@ from __future__ import annotations
 
 import gc
 import logging
-from typing import Any, List, Optional, Union
+from typing import Any
 
 import numpy as np
 import torch
-
 from vllm.config import CUDAGraphMode
 from vllm.forward_context import BatchDescriptor
 from vllm.multimodal.inputs import MultiModalKwargs
@@ -24,8 +23,9 @@ from vllm.v1.worker.gpu_model_runner import (
 )
 from vllm.v1.worker.ubatch_utils import UBatchSlices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
+
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.gpu.gpu_model_runner import OmniGPUModelRunner
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
 logger = logging.getLogger(__name__)
 
@@ -40,21 +40,20 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
 
     def _preprocess(
         self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-        ubatch_slices: Optional[UBatchSlices] = None,
-        num_tokens_after_padding: Optional[torch.Tensor] = None,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+        ubatch_slices: UBatchSlices | None = None,
+        num_tokens_after_padding: torch.Tensor | None = None,
     ) -> tuple[
         int,
         int,
-        Optional[torch.Tensor],
-        Optional[torch.Tensor],
-        Optional[torch.Tensor],
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
         torch.Tensor,
-        Optional[IntermediateTensors],
+        IntermediateTensors | None,
         dict[str, Any],
     ]:
-
         # Input token count for this iteration (not used by diffusion, but
         # retained to keep DP padding/ordering consistent)
         num_input_tokens = scheduler_output.total_num_scheduled_tokens
@@ -63,11 +62,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
 
         # _prepare_inputs may reorder the batch, so we must gather multi
         # modal outputs after that to ensure the correct order
-        if (
-            self.supports_mm_inputs
-            and get_pp_group().is_first_rank
-            and not self.model_config.is_encoder_decoder
-        ):
+        if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
             # Run the multimodal encoder if any.
             self._execute_mm_encoder(scheduler_output)
             mm_embeds = self._gather_mm_embeddings(scheduler_output)
@@ -102,11 +97,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             # If a batch only has token ids, then including the embedding layer
             # in the CUDA graph will be more performant (like in the else case
             # below).
-            token_ids_idx = (
-                self.is_token_ids.gpu[:num_input_tokens]
-                .nonzero(as_tuple=False)
-                .squeeze(1)
-            )
+            token_ids_idx = self.is_token_ids.gpu[:num_input_tokens].nonzero(as_tuple=False).squeeze(1)
             # Some tokens ids may need to become embeds
             if token_ids_idx.numel() > 0:
                 token_ids = self.input_ids.gpu[token_ids_idx]
@@ -136,10 +127,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
                 num_input_tokens, intermediate_tensors, True
             )
 
-        if (
-            self.model_config.is_encoder_decoder
-            and scheduler_output.scheduled_encoder_inputs
-        ):
+        if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
             encoder_inputs = self._extract_encoder_inputs(scheduler_output)
             model_kwargs.update(encoder_inputs)
 
@@ -157,9 +145,9 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
     @torch.inference_mode()
     def execute_model(
         self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-    ) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+    ) -> OmniModelRunnerOutput | IntermediateTensors:
         with record_function_or_nullcontext("Preprocess"):
             with self.synchronize_input_prep():
                 # Update persistent batch states.
@@ -213,7 +201,6 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             record_function_or_nullcontext("Forward"),
             self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
         ):
-
             outputs = self._run_diffusion(
                 input_ids=input_ids,
                 positions=positions,
@@ -225,24 +212,18 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
 
         _, multimodal_outputs = self.extract_multimodal_outputs(outputs)
         # Ensure one tensor per request, map to CPU for output struct
-        pooler_output: List[Optional[torch.Tensor]] = []
+        pooler_output: list[torch.Tensor | None] = []
         if isinstance(multimodal_outputs, torch.Tensor):
             # If model returned a single stacked tensor, split by requests
             assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
             for i in range(self.input_batch.num_reqs):
-                pooler_output.append(
-                    multimodal_outputs[i].detach().to("cpu").contiguous()
-                )
+                pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous())
         elif isinstance(multimodal_outputs, list):
             for out in multimodal_outputs:
-                pooler_output.append(
-                    out.detach().to("cpu").contiguous() if out is not None else None
-                )
+                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
         elif isinstance(multimodal_outputs, dict):
             for out in multimodal_outputs.values():
-                pooler_output.append(
-                    out.detach().to("cpu").contiguous() if out is not None else None
-                )
+                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
         else:
             raise RuntimeError("Unsupported diffusion output type")
 
@@ -272,11 +253,11 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
         *,
         input_ids: torch.Tensor,
         positions: torch.Tensor,
-        intermediate_tensors: Optional[IntermediateTensors],
-        inputs_embeds: Optional[torch.Tensor],
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None,
         multimodal_kwargs: dict,
         logits_indices: torch.Tensor,
-    ) -> Union[torch.Tensor, list[torch.Tensor]]:
+    ) -> torch.Tensor | list[torch.Tensor]:
         """Runs the diffusion process and returns per-request tensors.
 
         Tries model interfaces in the following order for maximal compatibility:
@@ -302,15 +283,15 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
         # TODO: add the diffuse method for other models
 
         raise RuntimeError(
-            "The loaded model does not expose diffusion interfaces 'sample', "
-            "'forward', or 'diffuse'. Please implement one of them or adapt the runner."
+            "The loaded model does not expose diffusion interfaces 'sample', 'forward', or 'diffuse'. "
+            "Please implement one of them or adapt the runner."
         )
 
     @torch.inference_mode()
     def _dummy_run(
         self,
         num_tokens: int,
-        cudagraph_runtime_mode: Optional[CUDAGraphMode] = None,
+        cudagraph_runtime_mode: CUDAGraphMode | None = None,
         force_attention: bool = False,
         uniform_decode: bool = False,
         skip_eplb: bool = False,
@@ -389,7 +370,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             num_tokens_across_dp = num_tokens_after_padding
             num_tokens_after_padding = int(num_tokens_after_padding[0].item())
 
-        attn_metadata: Optional[PerLayerAttnMetadata] = None
+        attn_metadata: PerLayerAttnMetadata | None = None
 
         # If force_attention is True, we always capture attention. Otherwise,
         # it only happens for cudagraph_runtime_mode=FULL.
@@ -405,41 +386,30 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
             self.query_start_loc.copy_to_gpu()
 
-            for kv_cache_group_id, kv_cache_group_spec in enumerate(
-                self.kv_cache_config.kv_cache_groups
-            ):
+            for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups):
                 common_attn_metadata = CommonAttentionMetadata(
                     query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
                     query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
                     seq_lens=self.seq_lens.gpu[:num_reqs],
                     seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
-                        :num_reqs
-                    ],
+                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],  # noqa: E501
                     num_reqs=num_reqs,
                     num_actual_tokens=num_tokens,
                     max_query_len=max_query_len,
                     max_seq_len=self.max_model_len,
-                    block_table_tensor=self.input_batch.block_table[
-                        kv_cache_group_id
-                    ].get_device_tensor(num_reqs),
-                    slot_mapping=self.input_batch.block_table[
-                        kv_cache_group_id
-                    ].slot_mapping.gpu[:num_tokens],
+                    block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs),
+                    slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens],
                     causal=True,
                 )
                 for attn_group in self.attn_groups[kv_cache_group_id]:
-
                     assert type(attn_metadata) is dict
-                    attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(  # noqa: E501
+                    attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(
                         common_attn_metadata
-                    )
+                    )  # noqa: E501
                     for layer_name in attn_group.layer_names:
                         attn_metadata[layer_name] = attn_metadata_i
 
-        with self.maybe_dummy_run_with_lora(
-            self.lora_config, num_scheduled_tokens, remove_lora
-        ):
+        with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora):
             model_kwargs = self._init_model_kwargs(num_tokens)
             if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
                 input_ids = None
@@ -465,17 +435,13 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
                 intermediate_tensors = None
             else:
                 if self.intermediate_tensors is None:
-                    self.intermediate_tensors = (
-                        self.model.make_empty_intermediate_tensors(
-                            batch_size=self.max_num_tokens,
-                            dtype=self.model_config.dtype,
-                            device=self.device,
-                        )
+                    self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
+                        batch_size=self.max_num_tokens,
+                        dtype=self.model_config.dtype,
+                        device=self.device,
                     )
 
-                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                    num_tokens, None, False
-                )
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens, None, False)
 
             # filter out the valid batch descriptor
             _cg_mode, batch_descriptor = (
@@ -491,23 +457,23 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             if cudagraph_runtime_mode is not None:
                 # we allow forcing NONE when the dispatcher disagrees to support
                 # warm ups for cudagraph capture
-                assert (
-                    cudagraph_runtime_mode == CUDAGraphMode.NONE
-                    or cudagraph_runtime_mode == _cg_mode
-                ), (
+                assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, (
                     f"Cudagraph runtime mode mismatch at dummy_run. "
                     f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
                 )
             else:
                 cudagraph_runtime_mode = _cg_mode
 
-            with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_tokens_after_padding,
-                num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_runtime_mode,
-                batch_descriptor=batch_descriptor,
+            with (
+                self.maybe_randomize_inputs(input_ids),
+                set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_tokens_after_padding,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    cudagraph_runtime_mode=cudagraph_runtime_mode,
+                    batch_descriptor=batch_descriptor,
+                ),
             ):
                 outputs = self.model(
                     input_ids=input_ids,
@@ -547,10 +513,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
         # Profile with multimodal encoder & encoder cache.
         if self.supports_mm_inputs:
             if self.model_config.multimodal_config.skip_mm_profiling:
-                logger.info(
-                    "Skipping memory profiling for multimodal encoder and "
-                    "encoder cache."
-                )
+                logger.info("Skipping memory profiling for multimodal encoder and encoder cache.")
             else:
                 mm_budget = self.mm_budget
                 assert mm_budget is not None
@@ -560,14 +523,11 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
                     # modality with the max possible input tokens even when
                     # it supports multiple.
                     dummy_modality = mm_budget.get_modality_with_max_tokens()
-                    max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[
-                        dummy_modality
-                    ]
+                    max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[dummy_modality]
 
                     logger.info(
-                        "Encoder cache will be initialized with a budget of "
-                        "%s tokens, and profiled with %s %s items of the "
-                        "maximum feature size.",
+                        "Encoder cache will be initialized with a budget "
+                        "of %s tokens, and profiled with %s %s items of the maximum feature size.",
                         encoder_budget,
                         max_mm_items_per_batch,
                         dummy_modality,
@@ -580,9 +540,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
                     )
 
                     # Run multimodal encoder.
-                    dummy_encoder_outputs = self.model.get_multimodal_embeddings(
-                        **batched_dummy_mm_inputs
-                    )
+                    dummy_encoder_outputs = self.model.get_multimodal_embeddings(**batched_dummy_mm_inputs)
 
                     sanity_check_mm_encoder_outputs(
                         dummy_encoder_outputs,
@@ -598,9 +556,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
                     if encoder_output_shape[0] < encoder_budget:
                         expanded_outputs = []
                         for output in dummy_encoder_outputs:
-                            expanded = output.new_zeros(
-                                (encoder_budget, encoder_output_shape[-1])
-                            )
+                            expanded = output.new_zeros((encoder_budget, encoder_output_shape[-1]))
                             num_tokens = output.shape[0]
                             expanded[:num_tokens].copy_(output)
                             expanded_outputs.append(expanded)

vllm_omni/worker/gpu_diffusion_worker.py

@@ -2,14 +2,14 @@ import gc
 import os
 
 import torch
-
 from vllm.model_executor import set_random_seed
 from vllm.platforms import current_platform
 from vllm.utils import GiB_bytes, MemorySnapshot
 from vllm.v1.utils import report_usage_stats
 from vllm.v1.worker.gpu_worker import Worker as GPUWorker
 from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-from vllm_omni.worker.gpu.gpu_diffusion_model_runner import GPUDiffusionModelRunner
+
+from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
 
 
 class GPUDiffusionWorker(GPUWorker):
@@ -43,10 +43,7 @@ class GPUDiffusionWorker(GPUWorker):
 
             # take current memory snapshot
             self.init_snapshot = MemorySnapshot()
-            self.requested_memory = (
-                self.init_snapshot.total_memory
-                * self.cache_config.gpu_memory_utilization
-            )
+            self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization
             if self.init_snapshot.free_memory < self.requested_memory:
 
                 def GiB(b):
@@ -65,9 +62,7 @@ class GPUDiffusionWorker(GPUWorker):
             raise RuntimeError(f"Not support device type: {self.device_config.device}")
 
         # Construct the model runner
-        self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(
-            self.vllm_config, self.device
-        )
+        self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(self.vllm_config, self.device)
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.

vllm_omni/worker/gpu_model_runner.py

@@ -1,8 +1,7 @@
-from typing import TYPE_CHECKING, List, Optional, Union, cast
+from typing import TYPE_CHECKING, Optional, Union, cast
 
 import numpy as np
 import torch
-
 from vllm.config import CUDAGraphMode
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import BatchDescriptor, set_forward_context
@@ -10,10 +9,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
 from vllm.sampling_params import SamplingType
 from vllm.utils import LazyLoader, cdiv
-from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata,
-    split_attn_metadata,
-)
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata, split_attn_metadata
 from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.worker.gpu_input_batch import CachedRequestState
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata
@@ -81,10 +77,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             sampling_params = new_req_data.sampling_params
             pooling_params = new_req_data.pooling_params
 
-            if (
-                sampling_params
-                and sampling_params.sampling_type == SamplingType.RANDOM_SEED
-            ):
+            if sampling_params and sampling_params.sampling_type == SamplingType.RANDOM_SEED:
                 generator = torch.Generator(device=self.device)
                 generator.manual_seed(sampling_params.seed)
             else:
@@ -145,9 +138,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                         if isinstance(payload_info, AdditionalInformationPayload):
                             for k, entry in payload_info.entries.items():
                                 if entry.tensor_data is not None:
-                                    dt = np.dtype(
-                                        getattr(entry, "tensor_dtype", "float32")
-                                    )
+                                    dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
                                     arr = np.frombuffer(entry.tensor_data, dtype=dt)
                                     arr = arr.reshape(entry.tensor_shape)
                                     info_dict[k] = torch.from_numpy(arr)
@@ -188,9 +179,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 new_token_ids = req_data.new_token_ids[i]
                 # Add the sampled token(s) from the previous step (if any).
                 # This doesn't include "unverified" tokens like spec tokens.
-                num_new_tokens = (
-                    num_computed_tokens + len(new_token_ids) - req_state.num_tokens
-                )
+                num_new_tokens = num_computed_tokens + len(new_token_ids) - req_state.num_tokens
                 if num_new_tokens == 1:
                     # Avoid slicing list in most common case.
                     req_state.output_token_ids.append(new_token_ids[-1])
@@ -228,23 +217,17 @@ class OmniGPUModelRunner(GPUModelRunner):
                 # Add new_token_ids to token_ids_cpu.
                 start_token_index = num_computed_tokens
                 end_token_index = num_computed_tokens + len(new_token_ids)
-                self.input_batch.token_ids_cpu[
-                    req_index, start_token_index:end_token_index
-                ] = new_token_ids
+                self.input_batch.token_ids_cpu[req_index, start_token_index:end_token_index] = new_token_ids
                 self.input_batch.num_tokens_no_spec[req_index] = end_token_index
                 self.input_batch.num_tokens[req_index] = end_token_index
 
             # Add spec_token_ids to token_ids_cpu.
-            spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
-                req_id, ()
-            )
+            spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id, ())
             if spec_token_ids:
                 num_spec_tokens = len(spec_token_ids)
                 start_index = self.input_batch.num_tokens_no_spec[req_index]
                 end_token_index = start_index + num_spec_tokens
-                self.input_batch.token_ids_cpu[
-                    req_index, start_index:end_token_index
-                ] = spec_token_ids
+                self.input_batch.token_ids_cpu[req_index, start_index:end_token_index] = spec_token_ids
                 # NOTE(woosuk): `num_tokens` here may include spec tokens.
                 self.input_batch.num_tokens[req_index] += num_spec_tokens
 
@@ -261,20 +244,15 @@ class OmniGPUModelRunner(GPUModelRunner):
         self.input_batch.refresh_metadata()
 
     @torch.inference_mode()
-    def extract_multimodal_outputs(
-        self, hidden_states: Union[torch.Tensor, List[torch.Tensor]]
-    ) -> dict:
-        if (
-            hasattr(self.model, "have_multimodal_outputs")
-            and self.model.have_multimodal_outputs
-        ):
+    def extract_multimodal_outputs(self, hidden_states: Union[torch.Tensor, list[torch.Tensor]]) -> dict:
+        if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs:
             text_hidden_states = hidden_states.text_hidden_states
             multimodal_outputs = hidden_states.multimodal_outputs
 
         elif isinstance(hidden_states, torch.Tensor):
             text_hidden_states = hidden_states
             multimodal_outputs = {}
-        elif isinstance(hidden_states, List):
+        elif isinstance(hidden_states, list):
             text_hidden_states = hidden_states[0]
             multimodal_outputs = {}
         else:
@@ -427,37 +405,25 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
             self.query_start_loc.copy_to_gpu()
 
-            for kv_cache_group_id, kv_cache_group_spec in enumerate(
-                self.kv_cache_config.kv_cache_groups
-            ):
+            for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups):
                 common_attn_metadata = CommonAttentionMetadata(
                     query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
                     query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
                     seq_lens=self.seq_lens.gpu[:num_reqs],
                     seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
-                        :num_reqs
-                    ],
+                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],  # noqa: E501
                     num_reqs=num_reqs,
                     num_actual_tokens=num_tokens,
                     max_query_len=max_query_len,
                     max_seq_len=self.max_model_len,
-                    block_table_tensor=self.input_batch.block_table[
-                        kv_cache_group_id
-                    ].get_device_tensor(num_reqs),
-                    slot_mapping=self.input_batch.block_table[
-                        kv_cache_group_id
-                    ].slot_mapping.gpu[:num_tokens],
+                    block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs),
+                    slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens],
                     causal=True,
                 )
                 for attn_group in self.attn_groups[kv_cache_group_id]:
                     if ubatch_slices is not None:
-                        common_attn_metadata_list = split_attn_metadata(
-                            ubatch_slices, common_attn_metadata
-                        )
-                        for ubid, common_attn_metadata in enumerate(
-                            common_attn_metadata_list
-                        ):
+                        common_attn_metadata_list = split_attn_metadata(ubatch_slices, common_attn_metadata)
+                        for ubid, common_attn_metadata in enumerate(common_attn_metadata_list):
                             assert common_attn_metadata.max_query_len == 1
                             attn_metadata_i = attn_group.get_metadata_builder(
                                 ubatch_id=ubid
@@ -467,15 +433,13 @@ class OmniGPUModelRunner(GPUModelRunner):
                                 attn_metadata[ubid][layer_name] = attn_metadata_i
                     else:
                         assert type(attn_metadata) is dict
-                        attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(  # noqa: E501
+                        attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(
                             common_attn_metadata
-                        )
+                        )  # noqa: E501
                         for layer_name in attn_group.layer_names:
                             attn_metadata[layer_name] = attn_metadata_i
 
-        with self.maybe_dummy_run_with_lora(
-            self.lora_config, num_scheduled_tokens, remove_lora
-        ):
+        with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora):
             model_kwargs = self._init_model_kwargs(num_tokens)
             if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
                 input_ids = None
@@ -501,17 +465,13 @@ class OmniGPUModelRunner(GPUModelRunner):
                 intermediate_tensors = None
             else:
                 if self.intermediate_tensors is None:
-                    self.intermediate_tensors = (
-                        self.model.make_empty_intermediate_tensors(
-                            batch_size=self.max_num_tokens,
-                            dtype=self.model_config.dtype,
-                            device=self.device,
-                        )
+                    self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
+                        batch_size=self.max_num_tokens,
+                        dtype=self.model_config.dtype,
+                        device=self.device,
                     )
 
-                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                    num_tokens, None, False
-                )
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens, None, False)
 
             # filter out the valid batch descriptor
             _cg_mode, batch_descriptor = (
@@ -527,10 +487,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             if cudagraph_runtime_mode is not None:
                 # we allow forcing NONE when the dispatcher disagrees to support
                 # warm ups for cudagraph capture
-                assert (
-                    cudagraph_runtime_mode == CUDAGraphMode.NONE
-                    or cudagraph_runtime_mode == _cg_mode
-                ), (
+                assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, (
                     f"Cudagraph runtime mode mismatch at dummy_run. "
                     f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
                 )
@@ -545,14 +502,17 @@ class OmniGPUModelRunner(GPUModelRunner):
                 if num_tokens_across_dp is not None:
                     num_tokens_across_dp[:] = num_tokens_after_padding
 
-            with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_tokens_after_padding,
-                num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_runtime_mode,
-                batch_descriptor=batch_descriptor,
-                ubatch_slices=ubatch_slices,
+            with (
+                self.maybe_randomize_inputs(input_ids),
+                set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_tokens_after_padding,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    cudagraph_runtime_mode=cudagraph_runtime_mode,
+                    batch_descriptor=batch_descriptor,
+                    ubatch_slices=ubatch_slices,
+                ),
             ):
                 outputs = self.model(
                     input_ids=input_ids,
@@ -582,7 +542,5 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.eplb_step(is_dummy=True, is_profile=is_profile)
 
         logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
-            hidden_states
-        )
+        hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
         return hidden_states, hidden_states[logit_indices]