revert doc deletion

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

.gitattributes

@@ -0,0 +1,44 @@
+# Set default behavior to automatically normalize line endings
+* text=auto
+
+# Explicitly declare text files you want to always be normalized and converted
+# to native line endings on checkout
+*.py text eol=lf
+*.yaml text eol=lf
+*.yml text eol=lf
+*.md text eol=lf
+*.txt text eol=lf
+*.json text eol=lf
+*.toml text eol=lf
+*.sh text eol=lf
+*.bat text eol=crlf
+*.cmd text eol=crlf
+
+# Declare files that will always have CRLF line endings on checkout
+*.bat text eol=crlf
+*.cmd text eol=crlf
+
+# Denote all files that are truly binary and should not be modified
+*.png binary
+*.jpg binary
+*.jpeg binary
+*.gif binary
+*.ico binary
+*.mov binary
+*.mp4 binary
+*.mp3 binary
+*.flv binary
+*.fla binary
+*.swf binary
+*.gz binary
+*.zip binary
+*.7z binary
+*.ttf binary
+*.eot binary
+*.woff binary
+*.woff2 binary
+*.pyc binary
+*.pdf binary
+*.ai binary
+*.psd binary
+*.xcf binary

.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/100-documentation.yml

@@ -0,0 +1,29 @@
+name: 📚 Documentation
+description: Report an issue related to documentation
+title: "[Doc]: "
+labels: ["documentation"]
+
+body:
+- type: textarea
+  attributes:
+    label: 📚 The doc issue
+    description: >
+      A clear and concise description of what content in the documentation is an issue.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Suggest a potential alternative/fix
+    description: >
+      Tell us how we could improve the documentation in this regard.
+- type: markdown
+  attributes:
+    value: >
+      Thanks for contributing 🎉!
+- type: checkboxes
+  id: askllm
+  attributes:
+    label: Before submitting a new issue...
+    options:
+      - label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
+        required: true

.github/ISSUE_TEMPLATE/200-installation.yml

@@ -0,0 +1,47 @@
+name: 🛠️ Installation
+description: Report an issue here when you hit errors during installation.
+title: "[Installation]: "
+labels: ["installation"]
+
+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/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/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
+      ```
+      It is suggested to download and execute the latest script, as vllm-omni might frequently update the diagnosis information needed for accurately and quickly responding to issues.
+    value: |
+      ```text
+      The output of `python collect_env.py`
+      ```
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: How you are installing vllm-omni
+    description: |
+      Paste the full command you are trying to execute.
+    value: |
+      ```sh
+      pip install -vvv vllm-omni
+      ```
+- type: markdown
+  attributes:
+    value: >
+      Thanks for contributing 🎉!
+- type: checkboxes
+  id: askllm
+  attributes:
+    label: Before submitting a new issue...
+    options:
+      - label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
+        required: true

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

@@ -0,0 +1,111 @@
+name: 🐛 Bug report
+description: Raise an issue here if you find a bug.
+title: "[Bug]: "
+labels: ["bug"]
+
+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/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+- type: markdown
+  attributes:
+    value: |
+      ⚠️ **SECURITY WARNING:** Please review any text you paste to ensure it does not contain sensitive information such as:
+      - API tokens or keys (e.g., Hugging Face tokens, OpenAI API keys)
+      - Passwords or authentication credentials
+      - Private URLs or endpoints
+      - Personal or confidential data
+
+      Consider redacting or replacing sensitive values with placeholders like `<YOUR_TOKEN_HERE>` when sharing configuration or code examples.
+- type: textarea
+  attributes:
+    label: Your current environment
+    description: |
+      Please run the following and paste the output below.
+      ```sh
+      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
+      ```
+      It is suggested to download and execute the latest script, as vllm-omni might frequently update the diagnosis information needed for accurately and quickly responding to issues.
+    value: |
+      <details>
+      <summary>The output of <code>python collect_env.py</code></summary>
+
+      ```text
+      Your output of `python collect_env.py` here
+      ```
+
+      </details>
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: 🐛 Describe the bug
+    description: |
+      Please provide a clear and concise description of what the bug is.
+
+      If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
+
+      ```python
+      from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
+
+      # Create stage configurations
+      ar_config = create_ar_stage_config(
+          stage_id=0,
+          model_path="Qwen/Qwen3-0.6B",
+          input_modalities=["text"],
+          output_modalities=["text"]
+      )
+
+      dit_config = create_dit_stage_config(
+          stage_id=1,
+          model_path="stabilityai/stable-diffusion-2-1",
+          input_modalities=["text"],
+          output_modalities=["image"]
+      )
+
+      # Initialize OmniLLM
+      omni = OmniLLM([ar_config, dit_config])
+
+      # Generate
+      outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
+      ```
+
+      If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.
+
+      Please also paste or describe the results you observe instead of the expected results. If you observe an error, please paste the error message including the **full** traceback of the exception. It may be relevant to wrap error messages in ```` ```triple quotes blocks``` ````.
+
+      Please set the environment variable `export VLLM_OMNI_LOGGING_LEVEL=DEBUG` to turn on more logging to help debugging potential issues.
+
+      If you experienced crashes or hangs, it would be helpful to run vllm-omni with `export VLLM_OMNI_TRACE_FUNCTION=1` . All the function calls in vllm-omni will be recorded. Inspect these log files, and tell which function crashes or hangs.
+    placeholder: |
+      A clear and concise description of what the bug is.
+
+      ```python
+      # Sample code to reproduce the problem
+      ```
+
+      ```
+      The error message you got, with the full traceback and the error logs with [dump_input.py:##] if present.
+      ```
+  validations:
+    required: true
+- type: markdown
+  attributes:
+    value: |
+      ⚠️ Please separate bugs of `transformers` implementation or usage from bugs of `vllm-omni`. If you think anything is wrong with the model's output:
+
+      - Try the counterpart of `transformers` first. If the error appears, please go to [their issues](https://github.com/huggingface/transformers/issues?q=is%3Aissue+is%3Aopen+sort%3Aupdated-desc).
+
+      - If the error only appears in vllm-omni, please provide the detailed script of how you run `transformers` and `vllm-omni`, also highlight the difference and what you expect.
+
+      Thanks for reporting 🙏!
+- type: checkboxes
+  id: askllm
+  attributes:
+    label: Before submitting a new issue...
+    options:
+      - label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
+        required: true

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

@@ -0,0 +1,38 @@
+name: 🚀 Feature request
+description: Submit a proposal/request for a new vllm-omni feature
+title: "[Feature]: "
+labels: ["feature request"]
+
+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/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+- type: textarea
+  attributes:
+    label: 🚀 The feature, motivation and pitch
+    description: >
+      A clear and concise description of the feature proposal. Please outline the motivation for the proposal. Is your feature request related to a specific problem? e.g., *"I'm working on X and would like Y to be possible"*. If this is related to another GitHub issue, please link here too.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Alternatives
+    description: >
+      A description of any alternative solutions or features you've considered, if any.
+- type: textarea
+  attributes:
+    label: Additional context
+    description: >
+      Add any other context or screenshots about the feature request.
+- type: markdown
+  attributes:
+    value: >
+      Thanks for contributing 🎉!
+- type: checkboxes
+  id: askllm
+  attributes:
+    label: Before submitting a new issue...
+    options:
+      - label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
+        required: true

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

@@ -0,0 +1,45 @@
+name: 🤗 Support request for a new model
+description: Submit a proposal/request for a new model support
+title: "[New Model]: "
+labels: ["new model"]
+
+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/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
+  attributes:
+    label: The model to consider.
+    description: >
+      A huggingface url, pointing to the model, e.g. https://huggingface.co/openai-community/gpt2 .
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: The closest model vllm-omni already supports.
+    description: >
+      Here is the list of models already supported by vllm-omni. Which model is the most similar to the model you want to add support for?
+- type: textarea
+  attributes:
+    label: What's your difficulty of supporting the model you want?
+    description: >
+      For example, any new operators or new architecture? Is it a multimodal model? Does it require special preprocessing?
+- type: textarea
+  attributes:
+    label: Use case and motivation
+    description: >
+      What is the primary use case for this model? Why is it important to support in vllm-omni?
+- type: markdown
+  attributes:
+    value: >
+      Thanks for contributing 🎉!
+- type: checkboxes
+  id: askllm
+  attributes:
+    label: Before submitting a new issue...
+    options:
+      - label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
+        required: true

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

@@ -0,0 +1,59 @@
+name: ⚡ Discussion on the performance of vllm-omni
+description: Submit a proposal/discussion about the performance of vllm-omni
+title: "[Performance]: "
+labels: ["performance"]
+
+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/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
+- type: textarea
+  attributes:
+    label: Proposal to improve performance
+    description: >
+      How do you plan to improve vllm-omni's performance?
+  validations:
+    required: false
+- type: textarea
+  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/vllm-project/vllm-omni/tree/main/tests/benchmarks .
+  validations:
+    required: false
+- type: textarea
+  attributes:
+    label: Misc discussion on performance
+    description: >
+      Anything about the performance.
+  validations:
+    required: false
+- type: textarea
+  attributes:
+    label: Your current environment (if you think it is necessary)
+    description: |
+      Please run the following and paste the output below.
+      ```sh
+      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
+      ```
+      It is suggested to download and execute the latest script, as vllm-omni might frequently update the diagnosis information needed for accurately and quickly responding to issues.
+    value: |
+      ```text
+      The output of `python collect_env.py`
+      ```
+  validations:
+    required: false
+- type: markdown
+  attributes:
+    value: >
+      Thanks for contributing 🎉!
+- type: checkboxes
+  id: askllm
+  attributes:
+    label: Before submitting a new issue...
+    options:
+      - label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
+        required: true

.github/ISSUE_TEMPLATE/750-RFC.yml

@@ -0,0 +1,52 @@
+name: 💬 Request for comments (RFC).
+description: Ask for feedback on major architectural changes or design choices.
+title: "[RFC]: "
+labels: ["RFC"]
+
+body:
+- type: markdown
+  attributes:
+    value: >
+      #### 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.
+    description: >
+      The motivation of the RFC.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Proposed Change.
+    description: >
+      The proposed change of the RFC.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Feedback Period.
+    description: >
+      The feedback period of the RFC. Usually at least one week.
+  validations:
+    required: false
+- type: textarea
+  attributes:
+    label: CC List.
+    description: >
+      The list of people you want to CC.
+  validations:
+    required: false
+- type: textarea
+  attributes:
+    label: Any Other Things.
+    description: >
+      Any other things you would like to mention.
+  validations:
+    required: false
+- type: checkboxes
+  id: askllm
+  attributes:
+    label: Before submitting a new issue...
+    options:
+      - label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
+        required: true

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,5 @@
+blank_issues_enabled: false
+contact_links:
+  - name: Questions
+    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

@@ -0,0 +1,21 @@
+<!-- markdownlint-disable -->
+PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
+
+## Purpose
+
+## Test Plan
+
+## Test Result
+
+---
+<details>
+<summary> Essential Elements of an Effective PR Description Checklist </summary>
+
+- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
+- [ ] The test plan, such as providing test command.
+- [ ] The test results, such as pasting the results comparison before and after, or e2e results
+- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
+- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
+</details>
+
+**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

.gitignore

@@ -0,0 +1,237 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+vllm/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# cursor
+.cursor/
+
+# docker
+docker/
+
+# scripts
+scripts/
+
+# tests
+tests/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be added to the global gitignore or merged into this project gitignore.  For a PyCharm
+#  project, it is recommended to include the following files:
+#  .idea/
+#  *.iml
+#  *.ipr
+#  *.iws
+#  .idea/
+
+# vLLM-omni specific
+# Model files and checkpoints
+checkpoints/
+*.bin
+*.safetensors
+*.pt
+*.pth
+
+# Cache directories
+cache/
+.cache/
+diffusion_cache/
+kv_cache/
+
+# Logs
+logs/
+*.log
+
+# Temporary files
+tmp/
+temp/
+.tmp/
+
+# IDE files
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS files
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db
+
+# Data files
+data/
+datasets/
+*.csv
+*.json
+*.jsonl
+*.parquet
+
+# Output files
+outputs/
+results/
+generated/
+
+# Configuration overrides
+configs/local.yaml
+configs/production.yaml
+configs/development.yaml
+
+# Docker
+.dockerignore
+Dockerfile.dev

.pre-commit-config.yaml

@@ -0,0 +1,32 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      # list of supported hooks: https://pre-commit.com/hooks.html
+      - id: check-yaml
+        args: ["--unsafe"]
+      - id: debug-statements
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+        args: ["--fix=lf"]
+      - id: trailing-whitespace
+        args: ["--markdown-linebreak-ext=md"]
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.14.0
+    hooks:
+    - id: ruff-check
+      args: [--output-format, github, --fix]
+    - id: ruff-format
+
+  - repo: https://github.com/crate-ci/typos
+    rev: v1.35.5
+    hooks:
+      - id: typos
+        # only for staged files
+
+  - repo: https://github.com/rhysd/actionlint
+    rev: v1.7.7
+    hooks:
+      - 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

CONTRIBUTING.md

@@ -0,0 +1,19 @@
+# Contributing to vLLM
+
+## Getting Started
+
+Setup environment with uv
+```bash
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+```
+
+Pre-commit has been linked to git commit. For run checks locally, please install pre-commit:
+```bash
+uv pip install pre-commit
+```
+
+Then run:
+```bash
+pre-commit run --show-diff-on-failure --color=always --all-files
+```

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,107 @@
+# vLLM-omni: Multi-modal Extension for vLLM
+
+vLLM-omni is designed to extend vLLM capabilities to support multi-modality model inference and serving, particularly focusing on non-autoregressive structures and non-textual outputs.
+
+## 🎯 Overview
+
+Traditional vLLM systems are limited to text-based, autoregressive generation. vLLM-omni addresses this limitation by enabling support for:
+
+- **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
+
+## 🏗️ Architecture
+
+vLLM-omni is built on a modular architecture that extends vLLM's core functionality:
+
+
+## 🚀 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
+
+## 📋 Supported Models
+
+### AR + Diffusion Transformer (DiT) Models
+- Qwen-omni (Thinker-Talker-Codec structure)
+- HunyunaImage 3.0 (Ongoing)
+- Qwen-Image (Ongoing)
+
+## 🛠️ Installation
+
+### Environment setup
+
+Use Docker to keep consistent basic environment (Optional, Recommended)
+```bash
+docker run --gpus all --ipc=host --network=host -v $source_dir:$container_dir --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
+```
+
+### Installation of vLLM (for users)
+
+Now we build it based on vLLM v0.11.0. Please install it with command below.
+```bash
+uv pip install vllm==0.11.0 --torch-backend=auto
+```
+
+### Installation of vLLM (for developers)
+
+If you want to check or debug with source code of vLLM, install stable release version of vllm with 0.11.0 from source with pre-built wheel file.
+
+```bash
+git clone https://github.com/vllm-project/vllm.git
+cd vllm
+git checkout v0.11.0
+```
+Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set VLLM_PRECOMPILED_WHEEL_LOCATION as your local absolute path of whl file.
+```bash
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.11.0/vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl
+```
+Install vllm with command below (If you have no existing PyTorch).
+```bash
+uv pip install --editable .
+```
+
+Install vllm with command below (If you already have PyTorch).
+```bash
+python use_existing_torch.py
+uv pip install -r requirements/build.txt
+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 successful.
+```bash
+python -c "import vllm._C"
+```
+
+### Installation of vLLM-omni
+Install additional requirements for 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)
+
+Please check the folder of [examples](examples)
+
+## Further details
+
+For detailed architecture design, see [vllm_omni_design.md](docs/architecture/vllm_omni_design.md) and [high_level_arch_design.md](docs/architecture/high_level_arch_design.md).

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

@@ -0,0 +1,143 @@
+# API Design Documentation
+
+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.
+
+## 📋 Module API
+
+### Core Module API
+**Core module** provides fundamental scheduling, caching, and resource management functionality.
+
+**Key Components:**
+- Request scheduling and prioritization
+- DiT cache management for diffusion models
+- Resource allocation and coordination
+- Inter-module communication
+
+### Engine Module API
+**Engine module** handles model loading, inference execution, and output processing.
+
+**Key Components:**
+- Model loading and initialization
+- Inference execution for AR and diffusion models
+- Input preprocessing for various modalities
+- Output postprocessing and formatting
+
+### Executor Module API
+**Executor module** coordinates and manages request execution across different engines and workers.
+
+**Key Components:**
+- Request routing and load balancing
+- Execution pipeline coordination
+- Worker management and task distribution
+- Error handling and recovery
+
+### Worker Module API
+**Worker module** provides the actual execution environment for model inference.
+
+**Key Components:**
+- Model execution and inference
+- GPU resource management
+- Request batching and processing
+- Performance optimization
+
+
+## 🔧 Implementation Guidelines
+
+### Code Standards
+- Use type hints for all methods
+- Include comprehensive docstrings
+- Follow PEP 8 style guidelines
+- Use async/await for I/O operations
+
+### Error Handling
+- Define specific exception types
+- Provide meaningful error messages
+- Include error recovery strategies
+- Log errors appropriately
+
+### Configuration
+- Use dataclasses for configuration
+- Provide sensible defaults
+- Include validation methods
+- Support environment variables
+
+### Testing
+- Write unit tests for all public methods
+- Include integration tests
+- Test error conditions
+- Validate configuration options
+
+## 🚀 Getting Started
+
+### For Developers
+
+1. **Choose a Module**
+   - Pick the module you want to work on
+   - Read the corresponding API documentation
+   - Understand the responsibilities and interfaces
+
+2. **Set Up Development Environment**
+   ```bash
+   # Clone the repository
+   git clone https://github.com/vllm-project/vllm-omni.git
+   cd vllm-omni
+
+   # Install dependencies
+   pip install -e ".[dev]"
+
+   # Set up pre-commit hooks
+   pre-commit install
+   ```
+
+3. **Start Implementation**
+   - Create the module directory
+   - Implement base classes first
+   - Add concrete implementations
+   - Write tests as you go
+
+### For Contributors
+
+1. **Review API Documentation**
+   - Read through all module APIs
+   - Understand the overall architecture
+   - Identify areas for improvement
+
+2. **Propose Changes**
+   - Create issues for API changes
+   - Discuss in pull requests
+   - Update documentation accordingly
+
+3. **Submit Contributions**
+   - Follow the coding standards
+   - Include tests for new features
+   - Update documentation
+   - Submit pull requests
+
+## 🤝 Contributing
+
+We welcome contributions to improve the API documentation and implementation. Please:
+
+1. **Follow the Template Structure**
+   - 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
+
+2. **Provide Working Examples**
+   - Test all code examples
+   - Include both basic and advanced usage
+   - Show integration patterns
+
+3. **Keep Documentation Updated**
+   - Update docs when APIs change
+   - Version control documentation changes
+   - Maintain backward compatibility notes
+
+## 📝 Notes
+
+- All API documentation should be kept up-to-date with implementation
+- Code examples should be tested and working
+- Configuration options should be validated
+- Error handling should be comprehensive
+- Performance considerations should be documented
+
+For questions or suggestions about the API documentation, please open an issue or start a discussion in the repository.

docs/contributing/design_documents/api_design_template.md

@@ -0,0 +1,336 @@
+# API Design Template for vLLM-omni Modules
+
+This template provides a standardized structure for designing APIs for core, engine, executor, and worker modules in vLLM-omni. Use this template to ensure consistency and completeness across all modules.
+
+## Template Structure
+
+### 1. Module Overview
+- **Purpose**: What this module does
+- **Responsibilities**: Key responsibilities of the module
+- **Dependencies**: What other modules this depends on
+- **Integration Points**: How it integrates with other modules
+
+### 2. Core Classes/Interfaces
+- **Base Classes**: Abstract base classes and interfaces
+- **Implementation Classes**: Concrete implementations
+- **Data Structures**: Key data structures and models
+
+### 3. Public API Methods
+- **Initialization**: Constructor and setup methods
+- **Core Operations**: Main functionality methods
+- **Configuration**: Configuration and parameter methods
+- **Lifecycle Management**: Start, stop, cleanup methods
+- **Monitoring**: Status, metrics, and debugging methods
+
+### 4. Configuration
+- **Configuration Classes**: Dataclasses or config objects
+- **Required Parameters**: Must-have configuration
+- **Optional Parameters**: Optional configuration with defaults
+- **Validation**: Parameter validation rules
+
+### 5. Error Handling
+- **Custom Exceptions**: Module-specific exceptions
+- **Error Codes**: Standardized error codes
+- **Recovery Strategies**: How to handle and recover from errors
+
+### 6. Examples
+- **Basic Usage**: Simple usage examples
+- **Advanced Usage**: Complex scenarios
+- **Integration Examples**: How to use with other modules
+
+---
+
+## Example: Core Module Template
+
+### 1. Module Overview
+
+**Purpose**: The core module provides fundamental scheduling and caching functionality for vLLM-omni.
+
+**Responsibilities**:
+- Request scheduling and prioritization
+- DiT cache management
+- Resource allocation and coordination
+- Inter-module communication
+
+**Dependencies**:
+- `vllm_omni.request` - Request handling
+- `vllm_omni.config` - Configuration management
+- `vllm_omni.utils` - Utility functions
+
+**Integration Points**:
+- Receives requests from entrypoints
+- Coordinates with engine modules
+- Manages worker allocation
+- Provides status to monitoring systems
+
+### 2. Core Classes/Interfaces
+
+```python
+from abc import ABC, abstractmethod
+from typing import List, Optional, Dict, Any
+from dataclasses import dataclass
+from enum import Enum
+
+class DiTCacheManager:
+    """Manages DiT cache for diffusion models."""
+
+    def __init__(self, config: DiTCacheConfig):
+        self.config = config
+        self.cache = {}
+        self.cache_stats = {}
+
+    async def get_cache(self, cache_key: str) -> Optional[Any]:
+        """Retrieve cached data."""
+        pass
+
+    async def set_cache(self, cache_key: str, data: Any) -> None:
+        """Store data in cache."""
+        pass
+
+    async def invalidate_cache(self, cache_key: str) -> None:
+        """Invalidate cached data."""
+        pass
+```
+
+### 3. Public API Methods
+
+#### Initialization
+```python
+class OmniScheduler():
+    def __init__(self, config: SchedulerConfig):
+        """Initialize the core scheduler."""
+        self.config = config
+        self.scheduler = self._create_scheduler()
+        self.cache_manager = DiTCacheManager(config.dit_cache_config)
+        self._running = False
+
+    def _create_scheduler(self) -> BaseScheduler:
+        """Factory method to create appropriate scheduler."""
+        if self.config.scheduler_type == SchedulerType.FIFO:
+            return FIFOScheduler(self.config)
+        elif self.config.scheduler_type == SchedulerType.PRIORITY:
+            return PriorityScheduler(self.config)
+        # ... other scheduler types
+```
+
+#### Core Operations
+```python
+    async def schedule(self, request: OmniRequest) -> bool:
+        """
+        Schedule a request for processing.
+
+        Args:
+            request: The request to schedule
+
+        Returns:
+            bool: True if successfully scheduled, False otherwise
+
+        Raises:
+            SchedulerError: If scheduling fails
+            QueueFullError: If queue is at capacity
+        """
+        pass
+```
+
+#### Configuration
+```python
+    def update_config(self, new_config: SchedulerConfig) -> None:
+        """Update scheduler configuration."""
+        pass
+
+    def get_config(self) -> SchedulerConfig:
+        """Get current configuration."""
+        pass
+```
+
+#### Lifecycle Management
+```python
+    async def start(self) -> None:
+        """Start the scheduler."""
+        self._running = True
+        # Start background tasks
+
+    async def stop(self) -> None:
+        """Stop the scheduler gracefully."""
+        self._running = False
+        # Cleanup resources
+
+    async def shutdown(self) -> None:
+        """Force shutdown the scheduler."""
+        # Immediate cleanup
+```
+
+#### Monitoring
+```python
+    def get_status(self) -> Dict[str, Any]:
+        """Get current scheduler status."""
+        return {
+            "running": self._running,
+            "queue_size": self.scheduler.queue_size(),
+            "processed_requests": self.scheduler.processed_count(),
+            "cache_hit_rate": self.cache_manager.get_hit_rate()
+        }
+
+    def get_metrics(self) -> Dict[str, float]:
+        """Get performance metrics."""
+        pass
+```
+
+### 4. Configuration
+
+```python
+@dataclass
+class OmniConfig:
+    """Configuration for the core module."""
+
+    # Scheduler configuration
+    scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)
+
+    # Cache configuration
+    dit_cache: DiTCacheConfig = field(default_factory=DiTCacheConfig)
+
+    # Resource limits
+    max_memory_gb: float = 16.0
+    max_gpu_utilization: float = 0.8
+
+    # Timeouts
+    request_timeout: int = 300
+    worker_timeout: int = 60
+
+    def validate(self) -> None:
+        """Validate configuration parameters."""
+        if self.max_memory_gb <= 0:
+            raise ValueError("max_memory_gb must be positive")
+        if not 0 < self.max_gpu_utilization <= 1:
+            raise ValueError("max_gpu_utilization must be between 0 and 1")
+```
+
+### 5. Error Handling
+
+```python
+class CoreModuleError(Exception):
+    """Base exception for core module errors."""
+    pass
+
+class SchedulerError(CoreModuleError):
+    """Scheduler-related errors."""
+    pass
+
+class QueueFullError(SchedulerError):
+    """Queue is at capacity."""
+    pass
+
+class CacheError(CoreModuleError):
+    """Cache-related errors."""
+    pass
+
+class ResourceError(CoreModuleError):
+    """Resource allocation errors."""
+    pass
+```
+
+### 6. Examples
+
+#### Basic Usage
+```python
+from vllm_omni.core import CoreScheduler, CoreModuleConfig
+from vllm_omni.request import create_text_request
+
+# Create configuration
+config = CoreModuleConfig(
+    scheduler=SchedulerConfig(scheduler_type=SchedulerType.FIFO),
+    max_memory_gb=8.0
+)
+
+# Initialize scheduler
+scheduler = CoreScheduler(config)
+await scheduler.start()
+
+# Create and schedule a request
+request = create_text_request(
+    request_id="req_001",
+    prompt="Hello, world!",
+    sampling_params=sampling_params
+)
+
+success = await scheduler.schedule_request(request)
+if success:
+    result = await scheduler.process_request(request)
+    print(f"Result: {result}")
+```
+
+#### Advanced Usage
+```python
+# Custom scheduler with priority handling
+config = CoreModuleConfig(
+    scheduler=SchedulerConfig(
+        scheduler_type=SchedulerType.PRIORITY,
+        priority_weights={"high": 2.0, "normal": 1.0, "low": 0.5}
+    )
+)
+
+scheduler = CoreScheduler(config)
+await scheduler.start()
+
+# Monitor scheduler status
+status = scheduler.get_status()
+print(f"Queue size: {status['queue_size']}")
+print(f"Cache hit rate: {status['cache_hit_rate']:.2%}")
+```
+
+---
+
+## Module-Specific Guidelines
+
+### Core Module
+- Focus on scheduling, caching, and resource management
+- Provide clear interfaces for other modules
+- Handle concurrency and thread safety
+- Implement comprehensive monitoring
+
+### Engine Module
+- Handle model loading and inference
+- Support both AR and diffusion models
+- Provide unified interface for different model types
+- Implement efficient memory management
+
+### Executor Module
+- Coordinate between different engines
+- Handle request routing and load balancing
+- Manage execution pipelines
+- Provide error recovery mechanisms
+
+### Worker Module
+- Handle actual model execution
+- Manage GPU resources
+- Implement batching strategies
+- Provide performance optimization
+
+---
+
+## Checklist for API Design
+
+- [ ] **Clear Purpose**: Module purpose is well-defined
+- [ ] **Complete Interface**: All public methods are documented
+- [ ] **Error Handling**: Comprehensive error handling strategy
+- [ ] **Configuration**: Flexible configuration system
+- [ ] **Examples**: Basic and advanced usage examples
+- [ ] **Type Hints**: All methods have proper type hints
+- [ ] **Documentation**: Comprehensive docstrings
+- [ ] **Testing**: Unit tests for all public methods
+- [ ] **Integration**: Clear integration points with other modules
+- [ ] **Performance**: Performance considerations documented
+
+---
+
+## Submission Guidelines
+
+1. Create a new file: `docs/api/[module_name]_api.md`
+2. Follow the template structure exactly
+3. Include all required sections
+4. Provide working code examples
+5. Ensure all methods have proper docstrings
+6. Include error handling strategies
+7. Submit for review before implementation
+
+This template ensures consistency and completeness across all vLLM-omni modules.

docs/contributing/design_documents/architecture_overview.md

@@ -0,0 +1,139 @@
+# Architecture Overview
+
+# Introduction
+
+This document outlines the architectural design for vLLM-omni, a new repository developed under the vLLM project. vLLM-omni supports multi-modality models with non-autoregressive structures and non-textual outputs, extending vLLM beyond traditional text-based, autoregressive generation.
+
+# Goals
+
+The primary goals for vLLM-omni are:
+
+* **Support for Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, alongside text.  
+* **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.  
+* **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.  
+* **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.
+
+# Expected Gains
+
+* It breaks through the limitation of vLLM that only supports autoregressive text generation models, and provides architectural support for popular model structures, especially multimodal generation.  
+* By supporting new output formats and new model architectures, developers can easily connect to mature ecosystems such as Higress Gateway and Ray, improving the ease of deployment and efficiency of corresponding models.  
+* Integrating efficient cache management, multiple parallel strategies, disaggregation from vLLM, to improve inference efficiency of corresponding models.
+
+# Non-Autoregressive Structure
+
+* **Diffusion Transformer:** initialized with a random latent noise, the transformer model iteratively predicts latent output for multiple steps, under control of sampling timestamps. Finally, it is decoded by the decoder of Variational AutoEncoder (VAE). It is not supported by vLLM now.  
+* **Autoregressive:** dominant generation paradigm for text, now supported by vLLM. It generates tokens conditioned on previous tokens one by one. The efficient KV Cache management provided by vLLM can accelerate inference effectively.  
+  ![][image1]                       ![][image2]
+
+
+|  | Autoregressive | DiT |
+| :---- | ----- | ----- |
+| Scenarios | Text generation | Multimodal generation |
+| Generation | Token-by-token, left to right | Continuous vector, iteratively refine |
+| Key advantages | Smooth, good local consistency | Global consistency, diversity and controlability |
+| Efficiency | More efficient with KV cache | Need generation repeatedly, slow |
+| Error accumulation | Easy to accumulate | Can fix possibly |
+| Key challenges | Limited global planning, long context generation | Large global computational cost |
+| Sequence length | Varied length | Fixed length |
+| Attention mask | Lower triangular | Full bi-directional matrix |
+| Parallel strategy | TP, DP, PP, .etc | DP, CP, .etc |
+| vLLM support | Supported | Unsupported |
+
+# Targeted models to support
+
+According to analysis for current popular open-source models, most of them have the combination of AR+DiT. Specifically, they can be further categorized into 3 types below:
+
+* **DiT as a main structure, with AR as text encoder (e.g.: Qwen-Image)**  
+  A popular model for image generation and editing. Main visual generation models are similar to the structure, such as Flux.  
+  ![][image3]  
+* **AR as a main structure, with DiT as multi-modal generator (e.g. BAGEL)**  
+  A unified multimodal comprehension and generation model. Visual generation can utilize CoT text generation output. Main unified multimodal models are similar to the structure.  
+      ![][image4]  
+* **Extension, multiple AR+DiT (e.g. Qwen-omni)**  
+  A new model for multimodal input and output. It designs the Thinker-Talker-Codec structure, which is a double-AR+DiT format.  
+  ![][image5]  
+
+
+
+# Architectural Overview
+
+vLLM-omni will import core modules from existing vLLM framework and introduce several new components and modifications to achieve its goals. The core architecture will be based on a modular design, allowing for independent development and integration of different modality handlers and output generators.  
+![][image6]
+
+## General ideas
+
+* It divides models in vLLM-omni into multiple stages, which are processed by different Engine Cores. Each stage can choose AR Engine Core for AR model part, and Diffusion Engine Core for diffusion model part.  
+* Designed a specialized diffusion engine core for DiT or dLLM.  
+* Plug-in with external DiT acceleration modules such as xDiT, Cache-DiT.  
+* Output linked with Higress API, Gradio for online demo, ComfyUI for less-code process.
+
+## Key Components
+
+* **Processor:** Input processing, tokenization or encoding  
+* **Scheduler:** for order and queue of incoming requests  
+* **Executor:**  Prepare model input and run  
+* **Multimodal output processor:** support multimodal outputs  
+* **Hidden State Output processor:** support intermediate hidden vectors as request output  
+* **Engine Core:** accelerate AR modules, imported from vLLM main repo  
+* **Diffusion Engine Core:** Wrap and accelerate DiT modules.  
+* **Diffusion model:** Wrap up diffusion models with vLLM format
+
+# Examples data flow of specific models
+
+1. Qwen-omni  
+   ![][image7]![][image8]  
+2. Qwen-Image  
+   ![][image9]      ![][image10]
+
+# Overview of function call dependency
+
+![][image11]
+
+# Roadmap
+
+* AsyncLLM:  
+- [ ] Initializer: resource allocation for different stages (2025 H2)  
+- [ ] Hidden state Output Processor (2025 H2)  
+- [ ] Multimodal Output Processor (2025 H2)  
+- [ ] Interface for API, Gradio, ComfyUI  
+* OmniEngine:  
+- [ ] Basic Diffusion Engine core: Diffusion step manager, Cache manager (2025 H2)  
+- [ ] Pipeline of different Engine Cores (2025 H2)  
+- [ ] Communication between Engine Cores via AsyncLLM (2025 H2)  
+* Model support (Included but not limited to):  
+- [ ] DiT: Qwen-Image (2025 H2), Flux-Kontext  
+- [ ] Unified-MLLM: BAGEL (2025 H2), OmniGen2, Janus-Pro  
+- [ ] Extensions of AR+DiT: Qwen-omni (2025 H2)  
+* Optimization  
+- [ ] Integration of other DiT acceleration modules, such as xDiT, Cache-DiT  
+- [ ] Disaggregation for AR and DiT  
+- [ ] Integration with vLLM's existing KV cache for potential future hybrid models.  
+- [ ] Support for streaming multi-modal inputs and outputs
+
+# Contact
+
+For any questions or feedback regarding this design document, please contact  
+Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk),  
+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
+
+[image2]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAARUAAAEOCAMAAACQIXOIAAADAFBMVEUAAAAAAAD/AAAAAAAAAADfACAAAAAAAADvADAAAAAAAAAAAADqACsAAADgACkAAAAAAAAAAAAAAAAAAAAAAAAtAAbmAC0AAAAAAADqADAAAAAAAAAAAAAAAAAAAADnADAAAAAAAAAAAADqAC4AAAAAAAAAAAAAAADpADAAAAAAAADoAC4AAADnADAAAAAAAAAAAACxACQAAAAAAADoADAQAALkAC8AAAAAAAAAAADoAC8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADqAC4AAAAAAAB/ABoAAADpAC8AAAAAAAAAAADqAC8AAAAAAAAAAAAAAADpAC8AAAAAAADpAC8AAAAAAAAAAAAAAAAAAAAAAAAAAADpAC8AAAA0AAvhAC0AAABXABLqAC8AAAAAAADpAC8AAAAAAADqAC+rACLpAC/IACgAAADGACjqADDpAC8AAADVACsAAAAdHRoAAAAdHRojHBowGhw2GR08GB1JAA9XFSBdFCFmABVpEiKQDSWWDCenACKq4+q8Biu9ACbCBSvTACvcAi7jAS/pAC/4Y4L9gpv+h6D/kKcAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAn++oGAAAAdHRSTlMAAQEECAgKCxATFRYYGRkaHSAhJCgoKCsvMDY3Ojw+QEJFR0hJTE1QUFRYWF1gYmRoaWttcHJyc3R3eHyAgYSGiYyNj5CSk5SXmZ2en6Klpq+vsrO3uLu/wMXMzc/S1tbX19fg4eft7u7v7/H09Pb3+f3+/iGfBl4AABcuSURBVHhe7Z1dcBNXlsePaeE4MW23jLyOhASOsPiwwYAXg5kwZAOJcajNkDDxUMkkVcmmaqt2HlJTNTOv+7Cvm62amtpK1W7NJFWbbDIDMyRVmYUAkxAIWcAQMAx2+Su2wY6FQTYysh0jS2bvR7fUfdV2d0stS2Lv70G6ffuclvqv27dvdx+dC8DhcDgcDofD4TwcFLEVOWNlcR9bZRHJlSjGB1T11nGwFblC2Op/m62zyMrmRDH4kKhiA/c6AMp8oVFUDLPrrPEwqTI4CFDvCx9j661TMKqIq93SdNdAFB1rK4pvxn114pEoSOt8UVKnh7DBI00PXSZF7NIiXhmgLsfm8UhQKKo81YhfA6N/moai5mVHdvgAWg97D+C64GHdndzwpIheA9Hr6LXoWRG7eIN/XPHCEoBlf7zPWmspFFWWD6MfunpvVfMnaMHxk+73wN3y3OOdF6D6aXdNJ2uNqG+Gtg4o2tIyR9ZiF3eLu+XxLuyyWs9DxRK2Ik85/fvuKAwehdJivBQ8PjZ2HbWco2Nj31yGNawxoqIJjp4ZGwudhC3EA7tcP4FcjiEXWGMwICkUVcbI61TESRr3FXLMxNrwawhKE1ZJvGVBOvy5V0lHMdjl7lzs4gO8QLWdn0I5ggBcFStcS8UZXIyRE29kKXm7qzZKUAM3aG8T8ksjILtMTRYTlzmVoS6Foor/Wdx34v0zy7hSqFDXmqNQVDkQ+7pnDCpal7IrskKBqCLApXNs3YKElSbiTDYa8xRIb1tE982LDyNThGAz7YSdMzeYVSYokLYSg331E0vW90bMHkHXhl/6p9hAdMn6ecZ4C1MgbQUugrdu/fmv2Or5GT805AjUrY+lI0oe3V8xQJDkMYt5kMuc/nmbw+FwOBxOIVEw45V/YyvU/IKtyJA8GPHbub8LbutXhvdVFAplxL+4cFX04KrowVXRg6uiB1dFD66KHlwVPbgqenBV9OCq6MFV0YOrogdXRQ+uih5cFT24KnpwVfTgqujBVdGDq6IHV0UProoeXBU9uCp6FKAqcdOPANMm/54zrxueZKu07IHP2aoENj1RzT+W/WwZW8WwZw9bYzf5dwRNdm9nqxgurDXSLVPyTxW002wNg7FumaLuV36lKueSx6bZGoYlUtjWLuJt8qchFXnYVmDG6EvNzZSwVfZi9AVywZzhPs+UZPd7Z3fraWK4z9luLEafnxOM9zn6aFa/OBsX96/Mck74m5d/ZzCSgxfvzT+Us8IvdNXVrcw1tw+9aTQi+bgsm0O5vFQFgjcNRyTnN7A1NpKfqqB9NvrTWPAmW2MjeapK8OYmtorlvJFuGZCnqsD5LUY7HTTULX3yVZVgyHCnDXVLn3xVxURjMdYtbfJWFePGYqxb2uStKsY7baxb2uSvKsY7bahb2uSvKvCxbytbpSX459ezJEseq2LcFoybU5rksyrGO22oW5rksypop9kaBmPd0iOvVQmG2BqWLDWWvFbF+Fon+H1WGkt+q2J8rXMpK40lv1UxvtbpykpjyXNVjPc5K40lz1Ux3ueu7/VSC2ZInqvS9dVBtorlgy3r2KqMyYEq9b/8qUHCQxVdQPd5AadLBhcGacA++UiT3Q0fKnn/6kme5tjoZBtOSpwxl7Z2sVUMXfarYk9bcfhhG1OzYu1r2po0URrLAlxiKzLGnrZSK4E3GUkQ/HwWwLdd3HlWbZMuJhrLOiMLq9jSVoQauFRSnVweHxsba/9deNV8PYElTDQW2w8hW1Qpr4z1KNl1E8TaaC7ajDHRmxrqZhFbvrhXRNdxlS42z6pDYsJzxNXrH1zpVpaEjQEHDH4j5/6rDrge3OgiSe+EljJ4cENZIZY43/j+iuIEwgHFabO7v7u60XF6xMRBZhE7VHFsgxvRsORnVHHR5LxJSI5ar7IHFa041eaK7b/GC9JeH3r1Crgnoiu82/+C84Bjp+KqsPcUdaIrV2w/0YnW+COPoPOdNILOQjb3LHaoskvqOwtHWpvGVTm6xcpnxXc0TUXY7w++h35jHxmXoaW+/yDvL3yCF2bel8/j2IyuaNl0OEqd1m394JfKSuwl7N+3BmcVb+qTp3v44FWwVRYbVKkIQA/AxB1xjayK+y3y1qk9fjz+2Cnc8Ef6/XSJZL2O9+16bNrjhy+UwQ02Iysu+t01ndQJtQXipHjF+/zklDdzQnay+xiyobf1ijOD+Juik7OqNvbZUfwm1BCKcfrdO+QmUpxk+Q7ISyTttRPk1N/qFUGcuFd2urSVrldWku4HYDihu4kTlRUybyvotDzhQe9LImIDHaAEUcvf2dgi4aXyPbiTiB0aARdMaZKqyi0KYtK0dpVcfjAuuRSn5PBV8SI9uepWXZetjSVzVcoroeoFWlzVpuxc/LTQsBWrMkFDkkI4raoatPQdLT3QrEo1I9BdxisTXlqubrGzw81clc2i/E3LRPXJ+YxfKo4qxwsirj0hhaEP95fKQim2xcRVZaSl4nT1aXmlykvD7BU7G0vGqqBLoPP0wHG8rj45x44f3MaM+OlFATriAO+w6hIhBE5HUgm5XFRBcmVTs9mAvFLlpcXWxpJxb+uWYv20hN43q/rbocvqJYBeoBcF6IgD3F+qLhGGIyUNKeVaCXXiihMUk5GzxkvL7BUbQ8IyVUVoVM4Zyf2WaS8h9xQURvqhGfXK0gGSJx0vkXh9AdWN90LTTlSW1tIyrvfvgPZpxQmkZY8oXsTph7hSy1Wnm61Km0yPIPTLy1NH4HOppAxZMONQU6tajH9ZKb7S4fQEr+A+Ai89/1w3OD0foqUzlb6mrbfKxfPocuC05P/5QMwrQudZ2Sk06vR89qIge/0cO8W+TW5YZvZK08dsXbqw/27IbrytWDquPjsvd96NJJaXl95SyqIzel8ektAl5LX09T8HqZ3aScNWn3VZSLxtyr8bMm0r1ohENItj6qzpqgWtGV1KNAWNk4arRg9gTZNpv7J4XHUZ9RuzySvrDCkcVVBjYatYrhrpZpbCUcVMYzHUzSQFpIqJxmKom0kKSBUTjcVYN3MUkirGjcVYN3MUkirGO22smzkKShXjnTbUzRwFpQpcGn2RrdIy+9FP7JClsFSBaysNdnryulFzMkOBqWK80xd8BrqZocBUgQsr2RqGyQ4j3UxQaKpMkkdnC3Ghqgy9rlvF1luh0FSBC+Tfqwvs9GRvI1r/BpYmbQpOlcntaJ9fXWinLwQ2vvoGW2mNxb2/Ygdrb2/ayNZp8MZfZausUniqzC68z+u2LqyZKdg7lPlN8WPKrzj1vWaFgpCYeXkeA11S7lAWVr9imP4rHpnncZE1CkwVw51+MB0xUs4EhaUK3mm2iuX+vRhbZZkCUwXttOE+xyYsdCn65F8WPSOW7sFP2T74K1uvpu4Zz8IGBhRcW4HZz/6LDUtMoeMjtsYa+DlloXFnYLnr2m22VsNUycTCBg8jS//ecKhmmBpqIbT9yr9olhj+OVk0a7eQYdLMlBGTd/OB5osnZ0W0Jz1n4fUrlOyeJQpVlezCVdGDq6IHV0UProoeXBU9uCp6cFX04KrowVXRg6uiB1dFD66KHlwVPbgqenBV9OCq6MFV0YOrogdXRQ+uih5cFT24KnpwVTgcjt2YiNQQfJWTcbby4UYbb1u/610ajlfffJT+mb+++TdzzcsOGcbR2APJhwXwG/3/ti8eWlWGS6qpGC6oIO9CDUQXMVLine truncated
+
+[image3]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAPwAAAEqCAYAAAAxuz3hAACAAElEQVR4Xuy9B3QUV57/246AMTkpC6Gcc84555xzQgkFQBIIRAaTsQkGg43JGWMwwQRhG2MbbGzjgMceezyzO7sz7+y+8z9n98y+s+/N9/3ubZUolVpCIAkEut9zvqdbVdW3qqvv5/5+v1vVLZVKSOgR9e3p43hBpcLoF1UYO2YMtKcbQl/bgmyNmbqOmKXvAjMjd1gYe8LazBv2VoFwsA6Cs00YXG0j4WEfA0+HOHg7JcLHOQn+rqkIcEtHkEcWQrxyEOqdg3DfAkT4FSLSvwhR5OigEsQGlyIupALxoZVICJtNrkJSRA2SI2uRElWHtJh67oy4RmTGNyE7cR53TtJ85CTOR25SKz22dDk7obnb86z4FmTGNSMjdj7SY+ZRm01IJSdHNCApvB4JoXWIC6tDbGgtYkJqEBVUhcjA2YgIqES4fxnC/EoR6luCEJ9iBHsXcQd55Xe6AAEeefB3z0GAeza930z4OafBl+zlnAovpxR4OibTeWFOIifCwy4BbnbxcLGNpXPHHANn6yg4WUXC0TIcjhbhsLMIga15MDkQNmYBsLUIoPWBoI/peeXnJiT0SLp7+gQHfsxLz2PSuHHQ1zGBkZ4dge4IE0NXDruliResTH2oA/rDwSoYjtYhcLONgrtdNLwcGewJ8HVJhp9rCgLdM7pgD/PJ47BH+jLYCxEVUIzoQAZ7GYEuwV6NxPBqDjsDPTV6TjfQsxLmdoGem9yMvOQW5KUwL6Dn980GAOlRPRi0EvitHHwGfSpBnxo9FymRaugTw+Zw6JnjybEEfWxwNQ1IlTQ4VSDcr5wDL0GvBr+QO8irEIGezLkEfi78CHx/l3T4kRn0PhL0nfZg8Nsnwd0+Ea4c+jhuJ+sYOFpFwZmgd7KMIOhDYU/Q21kEqW0ZQANDsABeaPD0bSfwr7z8PKZMmAQDHXOC3QnGBiyye8B8lidFdy+C3pt3QDvLIIo6YRx2Ft0Z7CyyM9j9OyN7sGd2t8gexeyvhj0miEX2cm457OrITrBH1yM9tkEz7AS6BHR+ysIuS9DL/2bbdAM+Vh3lGfDMLMpL0MfJgI8OmI1IGfDB3sUEOIvwDPpCsjraM/AZ7Mws0vu5UpR3zaKBrxN6p1RuL0cp2ifBw6En8E7WauCZHcxDYG8Z2gU8y6ZcBPBCgykG/EsvqPDqmJegNVUbM/VsYDrTjWBnoJNn+WBOZSOWtS2jNNMfh/cewNIFy6mDRlDUiqVOHE8dOh6+BHyAezqHXR3d8zvT+EJ6ZLCTA0opnS8lsMrJ6jSeOTmyjiCsQ3Pda1i9ZDt/nhU/9z7wiQz2ZhoE5uGTa7dx8J0zBDWDu60T7s4on0iwdz7PSVxAbbDUXp3S56e24PrlGzRwtCIxlFL6kDk06NQS6LWU2tfQ8VUjMqCKBiYC3p8BX4ni3Ga8vm4XNq3diezkeg5+CIHPonuQpxp4P48c+Lplw8clCz6u6TQApqrtSGaRniJ7kGcGDYapcGMR3lad1jM7WsbAwTKaR3cHy0isWrQKrvYRPLW3swgm+INo+1ABvNDgiQH/MtXv48e+BN3pehTd7dQ1+ywvbjMjL1x67zi+/PQyjzofX76A9pbF1NkT4e4QQ/VuFpKjizpr98zO6J7LozszAz01ppLS+RLkpVL9HFmNyoIFKMtt5ZE9L6WJIno9CtPnY9XiHWitX4uKvEWYXbgUORTVm2avRmV+OwFL0Zpq819//hUH9pxAWc5iAn4Bf6zMX8oje2X+Esyrfa1zAGhFddGyTvib0TZ3A3758Seq61soujdSPd+A0qxFdOz1vI7PjKesIoGl+jUEfgXBXYqOi9fw8ZWr+OLGDRpkjiAquAQRgaWICS2her8YkUH0vj1pcPPLo9cUwI/ef7hfDg1w2QR4GgfemyJ70+wWZCWWEeTxiAjIoOyIojvV8X5u8Qj1SyXwKboT9Ls2vU7Ah/Mob2/JgA8WwAsNru6+dwJjRlE6P340jHRnwdjQkafwViY+vHZf0tqOvPQCHNv3LjxdonD1/PuoKqunx7Norm/D9s3bMa9mIZa0rKYUNgUhBDyr3SP8iuixiOrmKpw59gHam9fh4N5jOLr/JBY1v4bjh89g+6a9+PLmF6gtXYKvb93htfvONw7gP/7+77RuH4F2C3t3HcVP399DTsICAnQx/vzrr3h313H88d7PuHn9C3R8eINeexenjpyn7X7C6qU78ObrR/DNnbu4crED1698ijtffoszxy/g+6/vUjSfQ4NOI76j52to23vf/YAFTevx6fVPsXThFpyi4wom2LMS5+L8+xeQGl+H/XuO4EZHBz66cg2b1mzDv//lD9izfQ+dgwvY99Z+Op49OH34KN7cvBOH9+ynQacNW9e+TmAnID6sBPt3vY20+CJsXLEBZXk12L/zLRRlVmFl2wpU5NdgYWMbwR2JhXMXwdlWAl4NvatNiABeaPDEIvyEV16A1qTRMDGwhKmRG6xMvQl2H2xdtwnnTx7HwuaF+ODkMaQl5GHjqg1U26Zg26atePvN3aitaMKurW+hJK8RLtbx2PXG20iPm40IqtmjKIUPdMvFLz99j2+/+grleQvw5edfoWF2O25+9BnBvZ/A/QP27jxMsH5P9fQc3Lj+Oa5e/Ag7Nu3HX//8Ozas2IUlzVv5JFzT7HW4ce0mDS6v4zOC/U+//Er7O4A///E3rF+5m2DcjZ9/+IkGkW/wze2vsap9F5Uja3Dl/DV8/833+OC9y0iMmIu0uPkUuT9B/ewVdBxf4K3th7Bo3ibUlC3H+uW76P2VoyCjFRcJ+KKsJnz31R1sXL0Tp4+cxm8/36VBaB9uf3qDBrwVlPl8hN1bd1PmUImFTStRlt+E8oJ5mFe9GB52yVRC1OG9w4eRGleG948dR0JkPna/8SYW1i9Cflo5qovrUVFQy2v5mJBMOFoR8BZhZAG80BCIAT9x7AvQnjwGZjNtYW7swWfkrU39CPIT1Klb4GIXhk2rN/NavrFqPuLCcwmQZWisXYgPz52nqH0I8+cspYiWhHWrXuepO0vhowLK4O+Si7//2+84/94FxIfOxu5tB/Htl9/g2qWPKOpvxq8//YrbN29ThP0ciSH1uPXpHRzZ+z7amjbhG4rcn16/hdfXvstr9tqS1fj9l9/x092f+Wt/+u5nitQ/4ucff8bZkx/S9t/jzq1v8da247h390dcufAp5lSsxK9/+JWWf43bn32LBErnU6Ln4S5lAN/c+Y4ymK3YuukgastWYnHzG1RSrKPjrkJMQCUNUnfw1c1blD28h8ykRhzaexz/8tsfsXTBOvzxx+8oU6jEiYMncLPjGhY0rMSyhWuQnjAbzXOWUQSvh4dDKjLiqnD60FEaQOZgdftayiAOU7SnMieqiAaC/Ti5fz+/0uHuEI/lC5ZT2h/ZDXgXWwG80CCKpfRTxr0EvWmvwtLYlUd3W/MAfi3Y3IiBHwAHijo25sE8+rjYxsCNOqe7QwI8HJPoMQmeTqm8fmWpfFxoGcK8C3h0T4upxe7tB7Fy8RYkhFUiIYRACpyNaIIpNqiKauk6xIfU8MieFM5m5ucjLbqJoqX6OjubcGP1N6vVi9IXkxdx8JnZ37lJbeSFXY/ZlPZnJ7QgN3ERvwafGcdm9Bchi+r2lOj5VK/PRWpsM79Ex1J75oQwquFD5lAdTw6uo+OrQWxgLT1W8dQ+MrAc65bt5DV+sHcZAjyKEOxTggA39fV4f8pg/F1z6P1nwd89m8/Y+7qyCbxMPpnn55zOr8fz2XoaEF1t2fV49hhP5zUGDhaRcLaOJUfDjkB3tIrgwLOJOwa8uCwnNKhiwE+b8DL0p42DpakHrM18OfBslpjdBGJnHkIdkzqiNcFuF8UjkYdjPLyck+HtkkKdms1AZyLII5sAyEGoF6vf2TX3Uh7RC7PmUx1bxaN9QiiDuxqJYbUEeC1SIhso0jVyyNlMegbByC6lsfSdAcwm4gpSF6MwjWDvcjvB3s4fC7lpfWo7ClIWIy+JnEwDQWIb1fyLqJ02Dj9rKyN+AdJiFxDsDP5mJEc1UYRuohS/iUf9OAZ8SB3iOqGPJrNIH+ZTSduyQaGKX58P8y1DODnMp5TeazGC+Yx9AQI98hHomYcA987LdDQAcPNLdQQ/m8CjgdHTIQXu7Jq8XSKfxHNml+dsorkZ7ErgnQTwQoOp786cwPSJo2A4YyJFcfXNNeoZYmniKJR3QlZjutrF8sju6UQRyzmFolgaRbcMPjsf5JGDEC92o00RvwzHgI8JKueODZpNwFQTVDU0CNTwyJ4SOYfMZsubCPi5ncDP74zM7G45Nruujt55yYtQQC4kqIsIbmYGeX4Km6knJ9Pz5CVqyMnZ8cyd0JOz4hcS7K3cXcBHzkdKxDwkRzTy7IIBz03QxwQR7IHVHPho6VKdXyV3ONX33N5q4ENYxPcg4Bn0nuwOPAn4bG4fgp3Z25ldskuDN4v0BLwHAe9G0d7FJq7zzrtYfo4d2fV4Swa9+nq8o3WQAF5o8MSAnzFpNAE/mWD34zfX2Fuqo7qDpdqsEzpTBHKzj+NpPAPe2yWVX3dmwLObTlhHD/ZUAx/pX4IIcnRQGcFTQVGzilJmAp4iewKP7vX8bjcJdgn49Bh2K6wG4JPuA1/Aonmn1bCzyN7OrQZ9scxq6DM7gWewq4FvuQ98uBr4eErtlcBLsHcHvoI7zJuVLqUI8SwiqyN8EIEf4M5uub0P/X3gM+DllN4NeHeW3tvEy4BnUT5KDbwli/JB/DZmlQBeaLD0/ZmTBPwYGGlPJ9j9qbNRJ7MKVQNPNbsEO7tRxN0hsVt093FRR3d/Sl0DGfBe+Qgl4EMpwkewSbvASl6zx1N0Z+awR8wh2Ot5Kn8fePY4D2lRc5HBoW+hqMzulOse5ZmlqM6e53aCnpvIovt90LPiWNLine truncated
+
+[image4]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAO0AAAFdCAYAAADv8OpOAABCCElEQVR4Xu2dB3gUxRvGlwRCh9BTCb2KVAWkKiDSiwKiCPxV7IACKlbABohBESk2pAkC0kR6SWjSW+i9Kr0TapL3P7N3c+x911I2l1z4fs/zPrmd/XZmd2fendkaTWMYhmEYhmEYhmEYhmEYhmEYhmEYhmEYhmEYhmEYJoUA8OtVqnqTj8rUfappcMmmOfyzNhXJKVJw1hxNWwSXbdo9vFK9IdUb59UYhjGHryrUfapFSIlnS2XPs6NOgdA9Imm3WSqcNfvuanmLrK0WGNxLYxgmZcyr1SL020qNhvlrmdCocDHMr9MR0fWfx+qGL7jVGqEV9Z/TtbJBF/1vVAPny8nYn6s1R818IagaWKhnEa1ITroeDMMkko/K1u1QPTDo4FslquNEszeBZz5AwtP9cbplT6FeLvVfi5641eZd3GjVD9fb9EVs63643fY9nG3Z2yFW6lLrdzC71tMQRUZny6zVoevBMEwi6BRSrmWpHHn3BmfLiXhh1GvCfBVzF5TGSpSWDWoDzOuL7Hn8sPLL9ujfqZpDDFWmTPrf2IjQnG9Pi6yVX2MYJvGUyJ63XW7/zIdr5w8Vpv1ADI076Mbq1KIEpg5vhCmRT1g03FGThz2Oazu7Aadfw8JfmuH6ru44uvJZTPu2kUOsUVPF/O8/fgz+fpkOly0V2puuE8Mw7mkrdLBx4QjsaPwSauYNwdLfWgBnegGHXsPtnS/jg1eroGeXiuhlot7u/hAqlgpEZj8NT7Ro+3vjlq3/R1eMYRgn+BtM+8/jXVExRyFsmN4WOPUmTq3qghfalHYY3pqtXLnzxGXLnn1f7wHfBmoMw7jHaNp1j3fTTbt+mjTtW3i1U3ndVN3efAf9Ph+GPoOGmqp+n32Nd78Yhi6vvomQsHCEFS/1xoABA9i4DOMOl6Y92wuBuTOjRr0ncPQOsP9aQop18AZ00fTj94C1B/9DoSJBl1s/+z++j8sw7vBk2gZPtdGNtufyvRTp0E1g75U47Lsaj/3X7fM7GAusP3YOgQUKXmrS8mk2LcO4w5Np6zdtnSLTHr4FLNq6Dy06dLadw/YdNATbTl+3M+26o2flvIuBgfl72q8hwzB2pMS0+8TQ9thd6L0onae08fgFNGvX0XJu3L40nn6yGDL5Z8ankaP0XpeaVohNyzDuSI5pd1+6hw3HzqPPwK9QskRxNG7eCgfEkPfA9QRi2jj0GfClbtj+r1bWbyHh8OsILxyArNlz68vsEUNmfXh8/LyMOyX0hv0aMgxjR3JMK3vXds93swx3C4Xpf4f9PFE3HzXt+18M0+e/9Exp/TYSTr6JIvkzI3uuQJtpZU+9at/xe8FhRd8SseXtVpBhGHuSYtoj4vx0+op1qFKjJrRSlaENnQNtTTy0OSd1YwYFh+gGlMNmW698ORa9P/5MzPfHIw8VQMVSeVAkJAwT/l6mX5jS870NLIs5dGPBgQNZ6foxDENIrGnlFd+f/pyPEqWsD1uMXgVtRSy0hRegLb8BrVUPPX3k73/aneMeEr3vplOXMOi70fDLkhu5Agvjt3lL7a4gy4tVy2IOxu5JSChAVo9hGEpiTLv3yj30/yrSYtZmYlg87QC0RRcthlWSxpXpIqZ4iZK6WdWFJil10Urek5U97K4Lt7FX3v6R92+tt3zEsleFmtutIMMw9ngy7RMt2uON9z5BtoAAaJ37iqHwKWHQ6/aG1XXRYtxGz+rGHTd3MQ6Tq8qyd5Xq/kZvVKxUCQOHj8bW/67i0C3b1ePrQi3s15BhGDvcmTZ/YFb4+flBKxgCrf/P0JZccWJWouViyPz7Lt245R96WDeu7EnlkPmN9z5G3pw5oFVpCG3A79AKhOhxE/5ernraK4ULF25tv4YMw1DaCR1q5My0eQMsQ+LBs6AtvuxoUKe6aDnXfayFvuy0Zf+Ic+K7mLlyoyWvnHmhTT8EbZnorUcshVa5HoIKFcJb73+MHNmyXQwLC3vTfvUYhrEjOCB72xz+/ocq5SmEnU1eRiWjacXwWCtUTBjsmjDtJScGdaMJ26HVbIpHHquLdz79AqVKl4UWUR7ax+OhLb1qOSdeIYbTU/dCy5YLosdHYGDgjR49ekyh68gwjAEMGOD3aGBwvWyZMs8pn7sA8vtntzftI80sJqOmTIyib0Gr2tDSwxYIhrY2wbn5ZZzozbOKofjSpUuRkJBQiq4nwzCECrkK1dUyaZeFEjbMaGeOaaV+3gDt4ToW4/60ztJr0xg5VJ59HFlFzOzZs6Vp36brxzCMEyZ81ahA+WKBH3dvVwZXdr6MgnmypNy0UtG3oY1ZDS0gK7TIBY7zlwojTz+InH7arSlTpqwUpq1K141hGBe0bhxeI1s2/7stGxRFzqx+0Go1T7lppeTFqS9nQPvziOM8q2kLZvW/2LZ1C/5mFMMklbLFC1fKlyvLy+LnXa1mM8vwVV44Si3J+77CtFrhkMs5m/D7tAyTXLJoWXNt1IKLWe69Lr5iue2TGlp5B9rwhdCCwo/6PdbieboiDMMkkor1Hn9MCys1TcscIM5Hs0PLki115JcZWbLm+Dc8vEAIXQeGYZJIoT7fNG3+ZBM8Xq8uGtatkypq0qBefL0XekTRshmGSSL37t1rGR8fvxxeQpT1XUJCQkm6HgyT4RENv49V7em8xCI8lEWYKIr4KlUR5V0UGkHXhWEyPAYTXBbG/VT8HCj/CtUSv/2E/KXEdFZnEvNyiV62lcrHm8TFxV0U61Dauh4B1r/+dBsZJkNBjeAKYeqrwiSXqcSsKzTWm4j1ui3XQ/y9JHTNmnZEGDiMbivDZAioCdI1wpPnb9DE+1w7fwZnrt6Vpo0Rpo2g28owGQLa8G2cmQGtwPP4S76Jng7YNPYl1KzUBDfpDAMbv2yCoGLNdNNOB4rSbWWYDAFt+IrYqCEo+MxP2HMpgc5KA+LQrHxpNB8UTWfYcWHtp6gZVAL32LRMRoY2fAvnMLNfQ7w5Zy9ui6mXH6uCRzoNRO0yRRFcpAiCitbFxpVj0aNhBRQuVBghD3cRtrIw5IV6KF00DEFBwQgNL4n6Lw615bpqTC80figcwSFhKN+gC16p/yhavT8HuHMKS0e+hvDQMISFh4rlyiLmYrx1qVvY/ctz6D7hMOThI+HKbvz+YRuEFyuGYkVD0azvROy8ZIm9c+Ef9K1REPPuxscMOMymZTIoNkcZubseg9sWx09bzuuTFXNp0HKVwLglG7Bu7rdoV8wPDz/eCj0+m4yZn7REsMjmrHTUjc3Im7cMnhk4GYdiVmDM241RwL8A/pOZJFxFxby5UbH9R4hZ/Qc+eroycmYLQqtv9yDqm2cQmisPJq3ciyNHtmHqhy3Q+puVOCu8GHdlHfrXCULUVcuqbR7fA6VzhWLnieM4tnspcmQqgI5jtuBigozdh8/r58PPV+JjfrzJpmUyKMqnRs4tGISGBTTs0y/6XBZGLI93l57T513a+Qc6BflhtfVU95/IjiiWORyx+lQCDh/Yi107tiFq8SyM7N8RZTLnwRVcx9/9G6L3zL24Y1kMuL0XBct1xZwT/0LzD0erj//Eju3bsV1oR8xi+Id2wsy9N7C036MoUKiyWgo7ZryLGnkz4/XIGdhw8LRIice9e3FIEKa9d3k7PqmdD7/cYNMyGRibGwysG/kCimsB0PvZu6uRv+LrWHHpnpi4g4OL3kc5/8LQ760IfnulBnLkeRLSsJd2z0eVMsUQHhaGsuXKoGRwLmgBFUUXuBlftQzB9Jj7l37jrq5GyZaRiPl3BzIF5EZIibIoWbKkrtJlyyLsiQ8QffwWvm4ahoIlnrUth4QziBrdE/kDNGQNfgT7r6iBuRhIn5mPl0rnxyIeHjMZmftuuM+g5kWgBXXUf1/88w00HroV1+Vp47WD+L1rBAJK3X+WokNpP4R3+Q1357+JMvlz3O9JD89Gt5IactV8TRh/A4a0K4rx2y39Me4dwYTuFdB16jncPbcdOav3wJwLasEE7Jo5BH/tkbeAgedLi3Pmpt/ov+PP7kSn93/FdlXIv7OghXbF3H2WsfOeMW1RJF8JxMXH71jA92mZjIq1+dvRpWwmZHt8EOTQc9WnT+DtJaeh97Mn1+DdatkR3KCvLbZGzmxoMiwatxb2Qrm8WXFKBgq++V9lZNH8UO7pwZBD7DWRHVD9peFYsikGP/ZriQjRUw7blSCvHiGrfwhaD1qAy6f3Y9XkfqhTJDPmH7EYUZo2tPlw/XfChS3Qcj+C92bs16dvn5qFLBV7Yckxy42gb5sXRb6yHeUtn50H2LRMRkVv7YSxg4di+vpj4tcprJwyFidiLU68emIVfo0chT9WHbFGXkXk8D+w+aTlBPfSrvlo07QJmrXvivUxe7Doh4/x1bgl1lhgxe8j8eWgz7A26js0L14IJ+9a0mNPrMPU7z5B15d7YsD307Dh5P27sUObRyC46lu26cOrfsfgPi+iYaMWeO71j3BJjY7jT6CgMOzEffJ6t/5U1E66rQyTIbC5wSSuXbqEaxbf2Di1diI+6NrdNn10THuE5iqCe4m4Bbykb20EFWtBk3H+0nXbbSbJqXn98NyoTbaHL+RjjXRbGSZDYGj3qcaFpV+gjhgOf/LzdPw1bwYeC82OQo/1oWFOOT2vJyqHVqTJ9pxZgj71y2IXecSRbivDZAjsm3kqEheLzevWYNWq1Th+4Rad65qEWzi6YbHLxxcTYk9g3cIFWHvQcuHKCN1WhskQ0IaeXnE5kpY3aF1At5VhmBQQpWmBUjSdYZh0iDRrtKZtk2LjMkw6x2pYGMXGZZh0iqGHtTOtTPtb0/IZYxmGSWOc9bBUbFyGSSe46WGpuMdlmLQmMT2sE23TF2YYxvsI07YVGigVrWlHox0NqnRUxUnZMmAYJu2IFh6OdjSrUtT9SIZh0gXRbFqG8S2i2bQM41Line truncated
+
+[image5]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAR8AAAEwCAYAAACOrXWuAAA5IElEQVR4Xu2dB5xU1fXHH21FsSCisrTV2LBFMAhG/EfFilEES2LDQiIREZKoUdSIGI1ExW6CGAEBFYMoTXrZoQlIbwssbQHpZWlL3z3/d97M2b1z3ntT9s3dmTdzfp/P77Mzt5z35s6737n3vrKGIRKlkHZt2wt+t6lK/HOJRKIUF+/IfrTARyTyoXhHdvLu7ftgz84DtvRUscBHJPKheEd2c+GO/ba0VLHARyTyoXhHdnJg0jRoddvtsGbleji496gtP9kW+IhEPhTvyE4e8r+hYBaFvMUr4ciBYlt+si3wEYl8KN6RyTjNKtwenGp9/+1wyMrKguVL8i34qHmpYIGPSORD8Y6M3rGlEIYOGQGjR461Xg/7bmQQPktXwYE9R2DQV4PNvHGwd9dBW91kWOAjEvlQvCOjt2zcAS2uuRZa33GX9fqHYaMt+KxYthp2bN0Dl15yqZV39GCJra6b9+4qgn279cBK4CMS+VC8I6M3b9gOJ5xQHc4551zr9ajhYyz4rFy2BrZt2mmt/2BeyVGw1eXGKdpOE1jjx0yC774dbsuPx3jKH83TBT4ikQ9FHRjXcMCEyeEDxRZwatQ4Gc4773xX+GBe8RGwAFRyJBxCCBuETvFhjHcctm/eBfUbNLDq4TQO141wW04gcTOeaZs8caq1D/y0v8BHJPKhDu8/bk2Htm7aBT+YkJk1Y27s8DEBgms/Y0dNCAMJwmHrzzthwtjJsH7NJgs4F154kVUPy82aMcfaFk7pOGRU7y88BIfM/cMp25gfJkDbu+6GwYO+M9MPQ9HeI1Y+xhP4iEQ+1EMPtoOB/QfBqhXrLDjc1PJmCxyxwMfq9ZUqWe/xCmha10GgLVu8ErKqZUH/fl9agLjookZWOQTHTTfebL1esmC5VQ/r0GjGOpNmGtM//7QvPPLwozB39kLo93l/qFa1GrzdoycUHy+BF194GT75sBeszd8g8BGJ/CgERFezI6/NX28BAS8mLA98cCSCp+Q/793PGung4vSpp5wKXw0cZINPq9t+G7puKB9mTp9j1ulrgigPdmwutEYz5Kef6gInVj/RHF2Nh/59B8JJJ54EPd96D0pKSqDZVc3hkXaPwdKFKwQ+IpFPBX97rmspfG5v9du44YPGNZxfXdm0FCqrlq+D0049zYTPNzb43N7qDut1ft5aePnFV6z6A8wREq4VNWncBDr88UlYtmgldPxTJ6hSuSqMGjHWBp8mja8EHLUJfEQi/yph8LmqaTPrdTzw+ftL3aw6A7/42oTPXrjil1fAH9s/AUsXrRD4iERprsTA5xjAxY0usfKWLFweM3z+9uwLVv3PPwtO1x5/tD289eY71tktgY9IlMaqV68+vP6PN83OXmAB4Zabb40bPvger3Zev3YTrF25AYr2HLHuA8P1mgH9v7LBB7eBr5eakNqycbsFPtwmTrtwoZkWobs+/xLkNDwHJo7LDYdPcQncdmsr6Nzpz7B8ySqBj0jkR21av80Eyi7rlPjdbe6GHv982zoFHit87v/9g/DgAw+VXgCIxsXidas3Wms30wIzYee2PaXwQbC81aMn3NP2XutMFZ7ZcrveB/cLt4+v+5gjIwQd7h+eTUNYYT4CS+AjEvlQCAoEwvbNu2HKpOmweH5ezNf5lBwHmDF1Nvw4bXYYNBAmeGHhnJnzrfrbt+yG88+/wKqH+TjimTp5hgUPDhzVwRHQIQtQP/04D9547U0ImPu4zxxl4Wn9PTuL5Dofkciv4tBAxwyf0FXKbiMXTEdI7DDB1vHJp+GO394Zth1ePpKxfOGOA471BD4ikQ/FOzIagZOoe7vIa/LXQ37eGlt6IizwEYl8KN6R0TgdwsXcl158xXo94vsfgne1L10Ne8wp0BN/+JOVd2j/EVvdZFjgIxL5ULwjk+k2B3xtf5jYAdvNncm0wEck8qF4Ryar6zJ2+LifoUqGBT4ikQ/FO7KTvx88DKpUqWJduSzPcBaJRAkR78hOxpHOgT2HU2q0o1rgIxKJRCKRSCQSiUQikUhUsco1jMfQPF0kEom0CaETMAxATzKMx3m+SCQSaVHAMEoCIfigcwVAIpFIp0IjnjDwkAVAIpFIi9SplpsFQCKRKKGKNOLhFgCJRCItynWBEBiGXI0sEon0SeAjEomSIoGPSCRKigQ+IpEoKRL4iESipEjgIxKJkiKBj0gkSooEPiKRKCkS+IhEoqRI4CMSiUQiUSqrQYO6IBanm81DuzI/1kUpJv6licXpYEPgk/riX5pYnA42BD6pL/ULg4AhFvvWAh+fSeAjThcLfHwmgY84XSzw8ZkEPuJ0scDHZxL4iNPFAh+fSeAjThcLfHwmgY84XSzw8ZkEPuJ0scDHZxL4iNPFAh+fSeAjThcLfHwmgY84XSzw8ZkEPuJ0scDHZ8oE+Gz+zoBVXxmw7ht7nhf3ed6AKy+0pyfb119hwKz/2NPT3QIfnyld4XNskoEHoKMHdbOXL491wWftIANqnWrAtqHB97gd/hnQWdUMODA2vO6h8cG8NtcG3/d1qUvG/X/29wZ0e9SAklz7vvjJAh+fKV3hc33jYOe6/BcGFIU6KALp1qtSHz4Ig0qGCYPQe9rOgs+DI7gRbxrQ+e4yeBwYE15fHeHhZ8Y66OUDgnXoPY0GEVgIMr+PlgQ+PlM6wodGCjw90dYFHwTP16+UvaftcMhYZSsFP2ssoxYaFRHUVOfUCQKIp/vJAh+fKd3gM+3jYAeLdXSDUxwsr9qpE459xz5d+aiLHT48HsZSQYK+oYl9m+ooB6dMKkwiwYemaDRtoukmvudlI8EH2w1B5rQNv1jg4zOlG3xo1LNtmD3PybVOCU7Nhr4RnIbQdO1rBV7UMdVyWAatwodAgGnqNEeFGcGpR4eyMjiFIiBY6y+PhQMiEnwQNq1blAGrvPDBPBz9+HnqJfDxmdINPndeE+xgtM4TyQgVp1EOjZ4IQAgeXo46vQofp3IIDkzHxWOCGB8Jkd0AEAk+aAQW5ZcXPhQH980tP9Ut8PGZMhk+z/zOPm1CUyd99bHge1yDcSrH13ywDi9HwEH4EIjcIMLPcvHtONVTIZgI+Pj5rJfAx2dKN/jQtGt2L3ueaqvTXlN2SprnYQzMU1/zcip81FP711xm96xeQUhgvlvnLg98qA6NpgQ+Ah/fKN3gg50PRz98+uNk7NQ1qtvTqZPStAtf83KO0y6Hcnx7NArieejyTLv42a7ywoc+j0y7RBWmdIMPmtZsln5hz1NNi788naBEkKAOzuvS4jKlOZVzqoMLyjwPrQJATXeDz8ddgttT96G88HEDn58s8PGZ0hE+aJp+kXHqQ6/V0/A49cI0HCk93qqsjNpBCRqY3qihAb+6MNjh+al2tVyHO4Om2FSGAIBp111hwOmnhG8PwclBwz8LGWNwkJQXPnglNMLTr1MutMDHZ0pX+KjG09nje0ZehMb7vxb2safzMlM+jByHPPqt4Kl0nk7GPLdYCAGnNSadRmhGmjL6wQIfnykT4OM3V/TVxjRa4tM9v1ng4zMJfFLTdLEhT9dhhJ3bIrifLPDxmQQ+qetEPwLEzZGmh36ywMdnEviI08UCH59J4CNOFwt8fCaBjzhdLPDxmQQ+4nSxwMdnEviI08UCH59J4CNOFwt8fCb1CxOL08WGwCf1xb80sTgdbAh8Ul/8SxOL08GGwEckEolEIpFI5D/xMwterDO2ZcOohHFLAkaJLc+LIRhX5CyzvYttbebFIFMnkSnbgeHBOmNb9jl8SnL17rct36NL91vgI9Ih24HhwTpjWxb4hFvgI/KzbAeGB+uMbVngE26Bj8jPsh0YHqwztmWBT7gFPiI/y3ZgeLDO2JYFPuEW+Ij8LNuB4cE6Y1sW+IRb4CPys2wHhgfrjG1Z4BNugY/Iz7IdGB6sM7ZlgU+4BT4iP8t2YHiwztiWBT7hFviI/CzbgeHBOmNbFviEW+Aj8rNsB4YH64xtWeATboGPyM+yHRgerDO2ZYFPuAU+Ij/LdmB4sM7YlgU+4Rb4iPws24HhwTpjWxb4hFvgI/KzbAeGB+uMbVngE26Bj8jPsh0YHqwztmWBT7gFPiI/y3ZgeLDO2JYFPuEW+IhEolgk8AlZ4CMSiUQikUgUr/j/1RGLxWJdNtT/28UzxWKxWJcFPmKxOCl2hU9pokgkEiVIAh+RSJQUCXxEIlFSJPARiURJkcDHX2pvep7pk3hGEoXHRyrtj6rhRnD/SFVD7yvktgtRZAl84ld1I3gAR2oXyufuoxYqhzIBPueaPmD6LMPeftw1QnXcVJHw6Wn6NZ4ocpfAJz7hwTvC9Fum15u+Pzy7VNhmCAlVWN5rR013+FxruMcrMMJBEosqEj6oEkNf7LSTwCc+4a/ybtPNTL9r2AFDcoIPggPTz2Tp8Sjd4YNt6na8FRjueW6qaPgUmG7OE0XOEvjEJxq9oGj65SQn+FDHihU+txrB8mSM1zn0V+3sCES13BHDPiLLDeWpVoVA5fntwkoExcvhvuBfDp9doXR1n0gEAIIxmvYXXw8NveYqMOz7jSMNvt/qVCwafPA9llH3D9uT7/8DSj7FVPefhCM3bJNo00GRYWOMwCeKeOfAzsg7OgrLEXyyjTKQLC4tEVk4tcPyaqemKYkKHyqnSk2jaSKuoVxaWqJMlM+BRp2UPoNbOT5NonJqG/E0is0hg+kIARxVOqnAsH9WDj0EvAr9aPDB1075fP/VOhST7z+KyqqwErlI4BO7aIShwoZ3NBId1KpxsZl3FjdhefXXGOUEAKdy6vSO4OAESBTlO3V42m+UWznqbLQ/buVwnxCAKKqDC8qqcCQZL3y41O2gIsGH2ukRJR/3H/eBT51whEX7SzH5/pOwrMAnBgl8Yhd2fBy5tFL8pBE8ELuXFbOEafQL/GbofffS3KCcpkvU8fLine truncated
+
+[image6]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAdAAAAHUCAYAAACKzHkWAACAAElEQVR4XuydB1wUW5q3Z77dmdndibuzOzM36DXca85iwAiIZDEAKqJiwpxzzllERUUlGlCM16yIIkGSmHPOCkoGQQWhn6+qoBXLcFFE6dvn7+//6+pTp9Ir3U+/p8459ZvfCL0hhIRKWNnZ2bx8+VJZzs3N5cWLF8p7uTwvL08pf/78ORqNRlmfnp5OZmam8l5ISOhtyZ8X+fMjS/4Mye+1r4XryGXy5ygjI0Ox9vP2qVLzQ++lDpCQkJCQUOmWDEctLGVAZmVlKa8yVOUfqNrXzy01P/Re6gAJCRVX8fHxDBgw4JX79++Pi4uLstyvXz/69OmjlMmW3/ft2/eN+sLCuugRI0aQmpqq/jiUiLTZpjYTlQEqt97Ikl+1y59ban7ovdQB0kXJf0zyH5Hc7Kdrfvr06aummF+LLl++zG9/+1v5w8Zv/+23b/n/vaNMWFhXLf89y3/rf/jDH7h//77641BikjNO+btP+/2nvU1SklLzQ++lDpAuKjQ0lI4dO9K5c2ecnJx0yl27dlXO+9ekK1euKAB13biAiznnOJ15Qlj4V+urXGLEvGH8+c9/5sGDB+qPQ4lJBqicgebk5CjNtV+iz4CaH3ovdYB0UcHBwdjY2CjNJ+oMr7T71q1bdOjQQX1JOi0tQGd5Ticq4RhhD48KC/9qHZsazcAp/b84QGXJ2eeXyDy1UvND76UOkC5KC9CSavcvSd25cwc7Ozt1sU5LAFRYn/w1AfqlpeaH3ksdIF2UFqBf8pfY59Lt27cFQIWFddgCoHosdYB0UQKgpUsCoML6ZAFQPZY6QLooAdDSpXcC9NFRjidHEpMkXBxHJ0YQ+iD4rS9x4a9nAVA9ljpAuigB0NKldwH0yK1DjJg9nDELRgt/oscuHCPFdIYCUfWXuPDXswCoHksdIF2UAGjp0rsAeujGQarWrEblimWoWa2C8Ee6Xq2f+Nc//k8BaYyUyau/xIW/ngVA9VjqAOmiPidAk0KXMNRtL68mwdLksX1qJ2ytLTA3N8fcypmDd9LIS7nF9c1jWRKRUXjzj5Y+AbRilSrMneLA7TPzuXl6nnARffvCQk4cnUK92j8xZr4AaGmzAKgeSx0gXdTnA2gq8xpWpYbLWM4nFky6nJfNINNvWHjgLklJSSRHe1Cr6zye3TvBqWnVaLs+4c1dfKT0CaAVqlTGdZajFGYvSPYULqqz/Lh9bj61qlcUAC2FFgDVY6kDpIv6bAC9tgCjgQcInt+Pqfsv5pdpshlsVRH/ywV1EnZT3nAgcfdPcGJ2Xez8E19t/inSN4AumumIRoJCbuIa4SKaDB8lExUALZ0WANVjqQOki/pcAA0d1pZF16SFuI00GhVQUKphYvPfU7ZKHerXr0/tCoZM33mWvIcniJlcUwD0HRIA/bwWAC3dFgDVY6kDpIv6PAB9zhKL+rTuP46xg2z4t6qjySxYM87oTxhaOdGjRw96uszlRhq8fBBLxIQaAqDvkADo57UAaOm2AKgeSx0gXdTnAOiz85uoZupIWFg44cei8OpanaGhGuKfahhsXYG1Z9+cqDn73nEiJgqAvkslDVBSvCR7okl6XaYpKMt7R33ZecleSn3tq1KW9P76Rbenck/3ffspfLxPtQBo6bYAqB5LHSBdVPEBmse+6daYj96NFpPJO50oY7qMVZEPGWD8f7gfS39ji+w7UYSNKoel35M3yj9WAqAfYQl4pMxniFFFatp0ZWP4YglOEjhfrGShWQVqm9lwN2WD9GvIR4KafAypfoavtLyGwJkmTN68jCPjqjLxgC+axJVsXuTIweilkLUWnkr1nkuvadL+Mvykfa5VgJybIME51Sd/XZZUJ2l1wfnIx10PcZMwr96UMKkeyd7528p1n8n78GL/gDqM3bYcTdqaTwapAGjptgCoHksdIF1U8QH6krP7fybibsrrouzb+PsFsCM2gaj967kYn/N6naTcjMfERQdw8Fr+U+E/Vb9GgN64caMEATqXYa1McRnUmJ6+EyQw+cGpfpj3bUU/q3Y8SliE18whHApfhCZhOfu8B3D+6gqOLLRiyvTuOLUsi2FvZ+Zvm8/pXSM4d8WNkOWd8VjsTL/e1uw+PJsL3m1x7NSW4FMSXJ96kxY9mvE9jRg+xZm4RAnQqTIoXVk6xJypY81pVNOUiBQJsremcCOgA10djRkyrBtxEmAPD63PxB3LJTALgP5aLQCqx1IHSBdVfIB+PX1tgEZFRSnPU/1cDg8PZ926dfIHjakrJ5cAQOcw2NgKD3czWjo5kpjlzfHZ7Vm03YWhZm15Ej8R2zo2zPGfDgluTO7Zit0xywieZ8Y0vzlsGlSVIf5LeXB9Ca4DWuMfPouxBj8wZMUM7qyz4O/1m2M9YxSn17THamBvsm4toK9pTdbtm8XOsY3oMGUEEVJmu2tQQ2ymjOXBvu40LGvAqcduuHczwqKNA9ceeRLQryoDNi4ickwDAdBfuQVA9VjqAOmiBEA/XXXq1GH48OGMGjXqs3n06NFYmFkwfM4wohM/P0AHtTDj5+hheHRrRMAJN6YOsef4yeEMMLWVADoJu4a2LNg0AxLdmNHPnL3HJYDObc30ADcOjq3BhL1evHjkjttgczZJAJ1pXZNNZzagOdeXzrY92H9zBZlH+uIwSso4Tw6jdydbHjxZT3rsCOyc7Ri3fhoTLVsRdN8D0ucysLmZlIFKWentedzbNwjf1QOZ37MOvX3nETVWAPTXbgFQPZY6QLooAdBPVyuTVpTEk+zX+a1jwOR+JQTQ1kTfXEBqYDsshndj9BAH7t4YT/9XAG2H65aZEOfKGMfm7It1/yBAF7Srzf5rW8g97YKdTTd+vrycp4ddsB/pTLwE5j5agB4fTofuEkA3TGOCRSsO3csH6KDm5sQ8dGV+t6aYmVmxedtoVgyo+8UBGh4X8laZcMlbAFSPpQ6QLuqjAZp8Etf5c5k7dy4+O6O58mrevqIr4WwQ2zcGka5JZe+2fdzPLOKxVfraADVtbcqDU0GsmJsfD9lu7h4cPH2fW+H++PwcTHrBpEyy0q8E4rk2kPTn+dDNOn+QhfPyt1vitQ/tneI1q9YoXyolBdATcW5w2YUKteszakI/Eh9NUACaEDcB0xotGL5aAmjCfBwb1eHAqeWvALpvdHVmBq+XYLaCxYPyATq/bW32Xt38ToCmnJ9Cd6umXH2yibjtDhh1cWRp5BJ8OzfHLXIVxE+mfa1mRF+fSddWjWk0aJJ09Ts5NKgevdbNJ/oLAVRuKs9/Yot4WsuXtgCoHksdIF3URwH0xgbqGdRk0ro97NmziwltTajWcgZx6nq/oFOzzKnwp9Yc1zzj8rnLpGYXosxH6EMAffbsmbros6uVqSlPbpzh8B45HnvYe+AQQ1r+D108b7LH5U/8ppw1gTde/8LY6vT/+M333bmWlMuj42uoXMmGTT/nb+vZrxF1Bi4jVarnWWIAnY1LQ2MiUj3JfuGK44/fM2WzG5rHY+jZxIK7GX5sG1gF40aVsbNqhpmUKe4/sYygmcaM3+DKmTm1+d7WgZFe01k5xJT1oTOZZVGdXVekDPRUb2xad2HbJQmgh6TlIV3Jlo57YEZrjBpXp0XzOvhHrObp07VknexLk5o1adOzGVY/mnA8dTmbhjbHosyPNG5hyqLRTbGfMYHDo+oxZmsJ9cJ9dJSI+DBOpMbgHehJ7ca1mb92LtEJ4mktX9ICoHosdYB0UUUHaDaetpXovOX1/LWZ57exc0hbdibB4cGGlKlUixZ9l3D16laa/LMyVSpXpUyF+niEPCYvL5WVg834oVJVrJsZ8be/dORa3i16tOnJsYQXnPWfSJ3vK/FjxXK4n0gmowhMVQM0JSVF6cXq6OjI//7v/5ZI82phmZiY8DRTO2WElJxv6kkV65XIfYsPjCxPjTZjcPeP5lyuVJAdSpfydjQd3p/Lj16wpHsdZuwr9IXxIo5B5ob4nkzEa3UJALTAeYW2l/eVpx3XWVAul2lhpayX18ljPuVxoIXGfmqXC5dpVMvafSjHS5L3+3qbwsfQHl8pS8hfVs6r4Ljqa/gYqwEqf2HLEN17cRdjFo6ikXFDpdPWv77/lxJr0ZT7ZS0AqsdSB0gXVXSAnse2ohVHlHbGNDZPGMGAYWMZM3sZUfEaDvc3oM7c/ElvT65zJzI7f6u4vRPpNHobT66tpm7t4TyVyoLGNaPh/1pyJe+2BNBehN26Ri+Diux7ApqTrlRq5MzAQ2naA79XMkC7dOmiLK9atYoGDRrw7//+78oX4o4dO1S1P79MpQz0aUY+QNNPbKBpo04cLxjyumd4JZwmzMdqiBsjj+WRsc0O28WH6e/izPn7DxloWpvttwrtjBy2jnJgxIYzeHt6lhhA9c2FATrRbQKhD47Qb6ILZcp/r/ydaL1400IuZJ/lVMZxybHCX8iX8y4wfM5Q/vKXv5Ce/uZ48V+b3qSHkJ4B9CoOlRsRoMx9kMOjK5eICZiLm93/MjRSQ8iAegzcW5CNnVtJjaqVqVypEuW++z86uYVyP6A/9R39kI9ybdcM2lfqLgFUzkD7cVjKWG2qWXH2mQTQ1Bt0tbfEbPWd14d+j+7evav0hJXn2ZX/O7SWP4z29vZ0796dbt26vdfy+o+1s7Oz4l69elG9enWyNXmk7B7EQGcpEz//OhvdPeRHxu85RR/HmfRfuJtFFsasO3mZrt2duPAgjfF2NdlyvaBynpwpZ+PV14IRARcEQD+jtQCtXeNHuvRz5H/+8T9v/K3I/v1//B4zu9ZYOJhJNhf+grZ2tKJK7cr87ne/w8HBQZnyU/sZ+1RrP6elLaMthA4hWeoA6aKKDlCIHGtIpYFb0d5dTAycxhSzPzMkIh+gQw/kSKUJ9P/ub7jdz69zcmVXnKbsI+GCH63Npynbxiy0o9m37QoA2oLine truncated
+
+[image7]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAQwAAAEbCAYAAAAvX7qwAAA00ElEQVR4Xu2dCXgUVbbHC5AgoIgITkIIwYERFB0RFxj1jeuMy3MBlTcgIouKKIIL7jMjqOMyioyCDqCCAgqIIIiyCzQgAsq+LwJhFRDClrAn59Wp7pPcPlXVqXR1JVWd8/u+/5fue2+dqq6u8+97by3RNEHwmH27D0IQpVOOfxZBEDyGJ2JQJIYhCKUAT0Su7D2H4MDeHNj/22FTXWlKDEMQSgGeiKrQLDZv2ApzZ82DxT8tgwO/5ZjalJbEMAShFOCJqOrA3lwITZ8DXbs8Dm+98Q4cOXjC1Ka0JIYhCKUAT0RVB/cdgTGjxsFlTS+Hdm3bw/GcPFOb0pIYhiCUAjwRScdyTkGu3qP4+qtv4IrLr4T27TrCidx8yDlwDI4ePgl7dx0wLVOSEsMQhFKAJyIawq/bfoOrr7oGxo7+Fr4bN9EwjA4PdILfdJN49ZXX4c7b74ITR/JNSWwlnAc5lH3EVO5WYhiCUArwRMw9eBx2bt0DVaueAYM+HgITvplkGEbH9p1g94690PnhLlC/fgPIPwGmJLYS9kT27Mw2lbuVGIYglAI8ERNpGAf35cKUid/DZZddYbw/lH3U1KYo4fAH4/ByMQxBKAWOHT5lDBl27dgH3+nmsH/PYUeGkacbxsRvp8DkCdOMYQclMl6vsXTRSuPsypaNO+CrkV+DvprI9Ry5xjpwyMMNgITtjurbhGdk8PqPlnfdDaNGjIHD+48ZZnZ4/1GjjRiGIJQCbe9rB8OGjIANazcbib12xQbYtX1vkYZhZGy5csYymNjYC8C/aEBdH30cMjPrwZBPPzcmTbENJvz2LbuM1yuWrDHaolHRBWH4F7Vx3RZ44P728GyPF+DUsTyoeFpFePvN3pB3Kh9efP7v8MH7/WHT+q1iGIJQGqRUTIEX9ETctH6LkczrVv5SbMPAX340hmmTZ8Bvv+6HJ7s/Deef3xC+GDYiyjB2bNltvF69fD3M++Fn+GTgYN08Vhtmgz0H1Mqla6Hy6ZXhkj9eApAHUKVyFej97z6Qn58PV17RDB5o18FoI4YhCKUDPPvMC3EbBgr04Qleq3Ff67aGGTz1RA9o2LCRbhgjLQ1j/epN8PcX/2ksO1TvhYwbMx4ubXIpdH6oi7EdFcqfBhdecKHJMC5t0hSwRySGIQilR0IMA9u0bXO/Y8P4x0svG8sO+2w4jB093uhRPNTpYWO4IYYhCP4FntCHEG4M40TuKXj6qWeMxMaJzqIMY62+jqGffQF33dkCpk2aYdyvMnXSdFjw4yJYtWyd0SazbiZAfrRhNNSHOffc3cqYAxHDEIRSID29Drz26huwcV1WOJnjmPTES8i3bNoB2zb/CrkHjkO3rt3h97+vD0OHfGFpGCuXroFft+0xTArXhZOdOAmKWrNivW4W9eCG626A/JP5hYaRlw+33HyrHvsJvc0GMQxBKA0wiXfv2Gf0DO5ucbeR9Hja04lhtP7bfXBfm7bGaU4STlwOGTwMnu3xPMwJzTPuRdFXY5gBToje0/JeY9iBJqGejiXhhV54WvdX3UjwqlM0pTdff9s4+4LmgtuKbcQwBKEUUK9/OHLohCEn12HghVt4dgTFkx7PeuToPYocvbcxauQY0PSkxyTHthjfyijsdCwnzzj9yu9dEcMQhFKAJyhd6Vmp0unwyUefmQzjoQc7Q7165zm60hOFvYL1qzcar4tjFEVJDEMQSgGeiGQY5cqVh4H9B0UZxq7tv0EH/W/t2ukAJ50ZBg5Fjuc6u1GtOBLDEIRSgCciCec09up/6fb2Du06GvMIO7fuhq2bdia0txCPxDAEoRTgiUgybkvfdzTqeRj4AB2a3OTtS1piGIJQCvBEVIWnS7+WJ24JgkDwRFSF8w/zf/gJev7zFRjw34+N05y8TWlJDEMQBEEQBEEQBEEQhEQxU9M68DJBEAQTIU3L1wUoViUIglBISDELUlQDQRAEBIchIWYWKL28Y1RDQRDKNiGLngVXQWNBEMouIQdmQYosIgiCYJiHySR0l5ArMwVBMBMSwxAEwSkhMQxBEJwSEsMQBMEpITEMQRCcEhLDEATBKSExDEEQnBISwxAEwSkhMQxBEARBsCEjozaIRH6XfqiW58euUArwL0Yk8qM0MQx/wL8YkciP0sQw/IH6pUBIE4l8o8y6Yhi+QwxD5FeJYfgQMQyRXyWG4UPEMER+lRiGDxHDEPlVYhg+RAxD5FeJYfgQMQyRXyWG4UPEMER+lRiGDxHDEPlVYhg+RAxD5FeJYfgQMQyRXyWG4UOSzTDWDDX+gVCUFvQ3tyuuBj2nQdPzzeVuNKefBuXKaZAzKRxf3earLtKgzY3mZVA1qoW3BZcbzJbjmtpbg7uu0SB/pjmO3yWG4UOSyTC63W1OGNSIl81tiysvDAMTv+rp4dfcMEgpFXVjmBy9XGaqBi0iJlCUYaCpoCm93N68fr9LDMOHJIth3HFVOEGGvmSuS4QSbRjYu6BeAr6n+PSe1K97YeLzGFxHp4bb5rNyNBU0jaD1MsQwfEgyGAYlCi+30/Q+4S4/6q3O5nrUzjHh+msvCb+3MoyT0zW4+YrYcbDXQ3E2jywsx97F7rGF7+0MA9Xjb+Gexvz/ht8PeFqDsf8yt7MzDIoRtF6GGIYPSQbDwGSjrn1RqnGmBhf/XoMNX4R1XROz2eCvMZZhUmIbfI1SDWPTiHDZkkHhNndeFU5qXv9m53A9zq307FCYzOW06MSOZRhoTBgLhyH4/s6rw8nPewyxDAPr1O0LgsQwfEgyGMbT/2f+9bcSJZ46J8DL6L0ajyYnqQzbYNJSAqvLURJjvZ2JYVuevE4Mg9Yfj2EUbJ9FnV8lhuFDksEwcP7CiWHYGQtOIuKvP74uSFw20agOSXr8X8RkWHKjqdAwg85+8HWhsPfBzcRrw0AFbR5DDMOHJINhYLLxX2wuo1dwVXSvgETlmGjqa74ONWH1XVcwD6Jqfn9zb4OruIZBw5vh/yxcvxiGUCokg2FQghZ1vYXdXAf2MIZHTr1iLwTbqBOSNAQhw8A4OAehtuFSextcVvMJsQyDTouqw53iGoYMSYSEkAyGgdI/ipHoKz8z15Hol5qXq8mP8xXYhgyElqOLpagNJrHahgvbv9zBXI6i5FXL7AzD6rRqPIZhZVJ+lxiGD0kWw0DpH6dAODyg1+qFWzjkwDK8irLjreHXnW6LjoPJjuWN6mpw2fnhRO7bPXr+g8wHk7DzHeF2akJS8mIZnlI9+8ywgRScJWFzHGgY6vaTcHluAPEYBl6LgVd88nI/SwzDhySTYZDwGgq8JDqXTVyqmvW+BksHmct5m1gxUHjKdOK/7dthvVUcNAi6WpMvk2hRD8luiORXiWH4kGQ0jKAIexlWk7CJlnpKOEgSw/AhYhilJ7z6kl/AlWjhMIVP4gZFYhg+RAxD5FeJYfgQMQyRXyWG4UPEMER+lRiGDxHDEPlVYhg+RAxD5FeJYfgQMQyRXyWG4UNUwxCJ/CpNDMMf8C9GJPKjNDEMf8C/GJHIj9LEMARBEIRSJT+k5fPJs3jFYwuC4ACeSG7kdVwxDG/g+yZugQwtkh7Tl+5CXscNkmHw9cUt0Mp5ErdnYXKb6uKVGEbyY/rSXcjruGIYWuLiimEI8WD60l3I67hiGFri4ophCPFg+tJdyOu4Yhha4uKKYQjxYPrSXcjruGIYWuLiimEI8WD60l3I67hiGFri4ophCPFg+tJdyOu4Yhha4uKKYQjxYPrSXcjruGIYWuLiimEI8WD60l3I67hiGFri4ophCPFg+tJdyOu4Yhha4uKKYQjxYPrSXcjruGIYWuLiimEI8WD60l3I67hiGFri4ophCPFg+tJdyOu4Yhha4uKKYQjxYPrSXcjruGIYWuLiimEI8WD60l3I67hiGFri4ophCPFg+tJdyOu4Yhha4uKKYQiCIAiCIAiCIJjgjxQXiUQiVZr6bwZ4pUgkEqmyNYyCQkEQyjS6H+SJYQiC4AgxDEEQHCOGIQiCY8QwiuZdXcd5oQLuHxK2GxpdXWw6aeFYVXhFKZDI7bhG1yIten9xYX1VWsCC07RwO/4+6qpKl+TrGssLhTBiGLE5XdcWLfog5fCDHjUoqkXxSFbD2KfrXM28r1T5wTB6a2HTSGTMpEEMIzbjtfABiabRmtURPKnQZHhZcUhGw6DeBSdLizaAoigJw0DwO5BehgViGLGhXz084O2GJVZJhWVXsjI70nTNjaiFZm8YXbRwm+ct6pAbtMI42EYFE2twpC6kWS+P20HrwO1AYm0HSq3DdQyMvKZtQc7TCnsXnCzNbBgYZ7IWvR1qXVGG8Y6uD7TCnooaj+8X/Cy0jkm6mil10suwQAzDHjSJHF21Iu+tkgfh5ZggTvcZrgPb4kHbQNcbkfdqTIyXrYWNC9tQr0c1JKxHQ8N6VF+ljraH1oHJjO/VHhNtx/JIG6vtQNTtaBSpp+2g5MU2IS2cjAjOAdntjywtug7NEj/HJVp4HbQNRFGGgcMJ1djpszfRwvG+YfX4frgW3uZWum5X6rK0aAMRNDGMWOCBrnZL8aCyGpaoSYW/0tgOE88JmPy8q84TlQyCwCTBMto2SprGBS0KobZ8HVSmrkN9j5CJUBlPVsRqO3hXHtvY9c6ytOiYf9ait4HMRu0t8H1BhvF65HUbpR4NQd0ebjBYb7XNCNapsQRNDCMW+Lmpd4Hgr5/a4yCwnap20dXGQXorE4Htea+FD0nwNU946v0g2J7XE5RwfB00z9Ir8t6qDaKWYyy+HnU7eDITOByhNpwszXoZBHswX2nhehrO8HXQezK3XkodgmV8EhWHGhSPDMVq6IG9lVd4YVlHDMMa6spaiSeNmlT0Xh0u2BkGP/gJ1TCoDYrmDUihSL36K8/hvRNC7Q3YbQeibgfFstoOxC5OcQyDfvGxR/Line truncated
+
+[image8]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAVwAAAGJCAYAAADG5YuwAACAAElEQVR4XuydBXgV19aG23tv7//f+u3f21svFHd3d5fgkgRCcHeXIIFAkKDB3d1dgwa3lhYo7m4tUOT797fCnJ6chJaE5JCTs77nWU8ye/bImVnzztpr75l5CyqVSqVyit5yLFCpVCpV7EiBq1KpVE6SAlelUqmcJAWuSqVSOUkKXJVKpXKSFLgqlUrlJClwVSqVyklS4KpUKpWTpMBVqVQqJ0mBq1KpVE6SAlelUqmcJAWuSqVSOUkKXJVKpXKSFLgqlUrlJClwVSqVyklS4KpUKpWTpMBVqVQqJ0mBq1KpVE6SAlelUqmcJAWuSqVSOUkKXJVKpXKSFLgqlUrlJClwVSqVyklS4KpUKpWTpMB1Ez1//hyzZ89G+vTp8d133yFlypSoWrUqLl686Fg1zmjJkiVInDix7G/r1q3lN/yZqlWrJnXtrXTp0o7VIojHJUOGDFI/VapUmD9/Pq5du+ZYLdaUMWNGpE2bFvfv33ecpYpnUuC6iaZOnYoECRLgrbfestk//vEP5MqVCzdu3HCs/sb1448/IlGiRLZ9/fDDDzFy5EjHauH07bffhvt9tKRJkzpWC6d58+ZFOC4JEyZEvnz58PDhQ8fqsaJ//vOf+Nvf/oY7d+44zlLFMylw3UATJkywwaRGjRro0qULatasKREdy/7973/j3r17jou9MS1YsABJkiSRffP29kabNm3k5kAonTlzxrG66ObNmwIuRsT8fZb9GaQZQVvHpVKlSlK/Vq1atrKvvvoKjx8/dlwsxjVmzBgEBwc7ZVuqNysFbjzXgwcP8MknnwhA5s6dG27ekydP8P3338u8Dh06YOXKlfJ/y5Ytbc3b/v37S1n+/PltyzECZBmb+RSb5O+//76U5cmTBz///LOt+U94spwwIzj5/7vvvovy5cvj9u3btnXaq0yZMlIvderUtjJfX18p27Ztm13NP7R161aZ37x5c+zdu1fsz5rohNvXX38ty/CGZK/Vq1fLtjmvR48eUpY7d26ZXrp0qS2S/u9//4vz589j//79tmMcGBhoW8/169flRsEoe/HixfK7WYfbPXjwoK3exx9/LMfv7t27tjJV/JQCN56LMOBFzmZzZLKASKDwgv/ggw+QLl06LFq0SOYToBZc2MS+cuWKTNMOHz4s0Zk1/fe//13+/u///q9EihTXwzJGp/Z1aFYdRzHFce7cOdmWJQKf6yDMI1NQUFC47dAYuR86dMixqujEiRNS54svvnCcJdq4caPMZ36VKlCggG29b7/9tu13MNVhRd/WfOumwN9BkFrz7Otwu7///rvU+89//iP7qsCN/1LgxnMxIuMFXrhwYcdZoj179sj8ZMmSyXSOHDkEDJ6ennj06JHAwILF8ePHbemJzz77TKJERsisY0W0zIlyPoHEnKQF3G+++caWoyTQWcZtvYp69uwp9QsWLPjSjrPixYtLHUbdY8eORf369WWa0Xhk2rx5s8zPnj274yzRsWPHbMs/ffrUBlwrqqesqHbYsGEy3aRJE5m2ImYLuCzj9ij+Ft68mP5gBEwpcN1HCtx4ru3bt8sFz86xyLRz506Zb6UMrIiVAN63b5/8X7duXfnLPCNHAjCiq1y5sqQk3nvvPZkXmbHZbAGXALfk4eEhQH7ZPtmrVatWsjwjTSsifFVZnWGMlh1l/TaO2ohMR44csR0HQt4CLlMVltKkSYN33nkHFy5ckOk+ffpInUmTJsm0BVymDCwtXLhQ0hX/8z//Y+usVOC6jxS48VzMMRIKbAZbOVNGYhMnTpSOJquTaOjQoTKPEGBdRrT16tWTeadOnRJwMO3AiI+Q2rBhgwCQwGWUy0432m+//SbljH6fPXtmA26jRo1s+8T87asAl5DmsnXq1JF1vUycR4hzP7h9S1YkffLkSbvaYeLvZE6Vv5XHyF6MxFu0aCHLcr2UBdyjR4/a6hGcPLYW0CMDLo+PfdqCKRwFrvtKgRvPxeisUKFCAgJe6JcuXZJOLk5b5ti0z5Ytm20eYUBVr17dVtaxY0dbj3qKFCkEztOmTZNpNueLFi2KEiVKSEoiOsC9deuWdLJxOXbavYoyZ84s9fv27SvTjNwZiX/66acvhXXFihVlGUbCVq6XoyA4Rpkg5jxGwpQCVxUTUuC6gWbOnGmDJWHHiNG+84pDqUJDQ231vby8bPOsHKefn5+tzH60A1MMLCN4CVVrBANBx5RDdIDL/bXG4FaoUAGNGzeW5WkHDhyQOuzU4vSmTZtkmjcBC54sz5s3r0wTqi/TihUrbL+JNyUuxwclrDKORrDSGApcVUxIgesmsi50CyY09rCzM4v/sxPHGkbFIUzWOFhrmNOOHTtsy9k32wlVC240RrvdunXD1atXZT7hzHL7EQnswGMZn65yFDuoPv/883D7aW+zZs2SetZwtQEDBtiWK1myZLi6HGHxV+OLV61aJSkS++V48yhVqlS4Bx+sCNp+1IM1PMwaG9y1a1eZZh6cYqcYp5m6sMT9t5azOs2slI8++BD/pcB1IzG9wIuaTXb+tfKtgwcPRtOmTW312AT/9ddfpZ4V4XFZ5oBfFoWxnPUdx75yeZZzfZY4NpjrigyG9vsYmVn7w3QFp/nXXtZ+RAVejsclshQEfxfnE+yWuC3+Dqs+Ac06VrrFOmb2+8L9t5azRlxwvv20Kv5KgatSqVROkgJXpVKpnCQFrkqlUjlJClyVSqVykhS4KpVK5SQpcFUqlcpJUuCqVCqVk6TAValUKidJgatSqVROkgJXpVKpnCQFrkqlUjlJClyVSqVykhS4KpVK5SQpcFUqlcpJUuCqVCqVk6TAValUKidJgatSqVROkgLXiJ+DuXbtWpw061M1KpXK9aXANfrXv/6Fd/7xT/zj7/+IO/a3f+Cdv78jHyFUqVTxQwpco0/+/QmWHVmMvfdC44wd+HUPNp9bj/f+V4EbX8QPbNauXdtl7JNPPnH8CarXlAIXYcCds3MmtlzYGGds6+XNchN4/1/vO+6uykXFT7+7kpImTepYpHpNKXChwFU5RwpclQIXClyVc6TAVSlwocBVOUcKXJUCF+GBu+vmdoTe3IHQWzuw68a2CCB0lilw458UuCoFLv4AbsilTchdKBcy58iEDFkzoGTlkthxfauUb720OQyE5q/1v9jFsLKQi5sinW8/zTr2da35XL/Mt9+OAjfeSYGrUuAiPHDzl8iH/Q9244enh9BnbC8MnDEA7fu3hVeTmth2dQtqNfeGb2sfiYQJxfl75qK8ZzlMWjcB265sQZ1WPmjQoR52394p8PRtUwc1GlbH9qshmL19BsrXLIvBMwOx49pWAatn4xroMqSTLNt5cCczXRMbzqzFtstbIgXumDFjsHTp0nBlKtfQXwH38OHDSJEiBVq1aoXOnTsjffr0ePDggWO119a9e/fQtm1bx+IIUuDGvBS4CA9cRrf9Jvqj1xg/lKlWGnNDZ5toN73As7V/SwTNGSwgbtW7JTadW48iZQtj373daBvQBu0D2mLapskYNCMQ7cz0kNmDMGrJSCw9vAhjVwQjb9Hc+OHZYfQZ10tAXaVOZcwImQafFrUQbOplz5sN60+tCRfhvvcCuIsXL5Zxke+++67D3qtcRX8F3KNHj6Jdu3a26caNG+PHH3/EtGnTkDlzZlSsWFHKL168iGTJksHT0xN169bFpEmTsGnTJuzYsUNA/fjxY9SsWRNZs2bFyJEjcfPmTWTKlAnp0qXDDz/8gJ49e+LLL7/E9evXbduKTArcmJcCF+GBm6tgTizYO1ds5Y/LsOn8BhQuUwg7r2+DdxNPgeea4yslch27YjQqeHlItLrv/m54mig4VfrUSJc5nUS6y39YigIl8iN3odySRhi1eASSp0oGv5HdMXrZKOQpkhvrTq7G8AVD0SWoE0pXLhkupbDcAPfL/34poH3rrbfEeNE2atRILsYmTZqgefPmaNmyJdq0McBv3x4dO3ZEly5d0L17d/To0QO9evWCv78/AgICEBgYiEGDBiEoKAjDhw/HiBEjEBwcLFHz+PHj5cKdOnUqpk+fjpkzZ2L27NmYN28eFi5cKMBftmwZVq5ciTVr1mDt2rXYsGGDXOghISHYvn07du7cid27d2Pfvn3Yv38/Dh06JBAhNH7++WecPHkSp0+fxrlz53D+/HlcunQJV65cwY0bNwQK/Buf7fvvv3d0vXA6duwYPvroIyRMmBBfffWVnFsqQ4YMEpXyYQQe4759++LgwYO4fPkyGjZsKOeL54DHvnfv3vK3WbNmuH//PgoUKCDnjOf49u3bePLkiRzrTp06OWw9ohS4MS8FLsIDN2/RPNhsIGvlWDeeDYtit18LMVFtc4xYOBQBE/uiVa8WWHNiFfIWyYODv+1Dsx5N0K5/W8zdNdtExbPQZ0xPDJ8/FGtNnWmbJqHjwPZo2bM5TuMkvJt5Yf7uOahsItzZO2bCt3UdjFgwDCUrFRfQ2ke4H77/IUJDQ5E/f34BboIECeQdC7zYaIx2Lly4IBA7e/asAO2XX34RuJ04cUJAxwuZ0CP8jhw5IiA8cOCAGOG4d+9euUi5HV7QhOe2bduwdetWbNmyBZs3b8bGjRsFsOvWrRPg0latWiUX8/LlyyXNsWTJEixatEgAvWDBAsyfPx9z584VcBPgM2bMkGhtypQpmDx5sgB+4sSJmDBhgoB+3Lhx8dq+++47R9cLJ54f3jQpRqKEKcXolseIx5jnkhEqzysBypsvgctztn79egEuz2OdOnXkJsrzRvF8MXrm+eWNztrOn0mBG/NS4MIOuCYKbd2nZTjgbjq3QXK4jDy3XwlBtXpVJe/KEQzMu07eMAnFKxTDhDXjJO/q2agGKvlUlBzuZrNs1bpVJP/LacKadRnpMirm/MqmbruAtrKujoHtbR1okXWaEYDM8RGYKtLine truncated
+
+[image9]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAOsAAAEWCAYAAACHcWCmAAB5cklEQVR4XuydB3gVVfrG4/5XRXonjYQQ0nvvvffee++kQwKkQOgdpIogiPQmIBBAJXRUFERQV1111+7aV1d33fX9f2duJlwmNw0IIXDe53lz751y5szM+c33nZlzb5SUuLgUSH3cSAz4sxIe//MjUBmrifEq+mQjTFCzwCRNW+hq2cFgkiNM9F1hZugBC2Nv2JgEwM40CA7moXCyjICzVSTcbGLgYZcAT/tEeDsmw9c5FX4u6Qhwy0QgOcg9CyFeOeRchPkUINy3kFyESP8SRAVMRnRgKWKDyxEXUoH40EokhlcjKWKK4OSIqUiJnCY4OaJWcFJ4zS3vE0JrEB8yVXB0UBViAqsQHVCJSL9yRPiWIdS3FCE+kxHsXYJAzyIEeBTC370Afm658HHJFuzllNnqDHg6ppHT4W6fCje7ZLjbJcHdNgGuVrFwIjtaxcDRIpqOQRQcLMjmEbA3DYetaRisTULIwbAyJhsFwtIwABYGfjDX94GpvjdM9DxgrOtOr+4wM/AAnYY/Sc8LF1c7qY8dgSce/RMGPvZ/GK86CVrqppg4noFqA72J9gKoxnoMVHeYG3pRA/RrBTWEQA2Hi3UUXG2iCdR4eDkktYHq75KBAGYGqkc2gj1zCNJ8weG+xYjwKxZAZZDGBJUJoDJIE8KqBFCTIwnQqBrBqdG15OkyR8nMwBVfZRBPR2JYLUE7BXHBU6jMagK2GlH+FYLDfUoFh5FDCNhgjyK6iBTA3zWfQL0JrLdzlmAZsBnwcGBOIWhT4ErQulnHwZXswqC1jIGDZbRgewasWSTszCJgQ8AyW5uEwtIoGBaGgbAiYC0N/GEhAOvZZgsjTw4rV/ekNno4Bg34Pwx54lFoqhlCW8MaOhPsCFQH6Gs7wUjXRYgAFkbe8HKJkIFqFkxRJUSIqK4UUd1s44SI6uOUQqCmUESVQSoDNYtAzaOIWohQ7wKCpVAAVYimAaUEZhUyE6ZRRKVoGiYXSaOYa5ERV4fJ2XOQHlNHnoG0aNF15OnC+9QoGaxJ4bVURg1iKbrmpTYgO7GOQK0iV9IFokwANZRFV69iqlchRft8BLgWEKQ58CZQ/d1zCVCKrI4M2NYIy0ywutkxJ9P+JsDZmlkGrDMB60wR1p4irL15pGBbs3CCNRTWxgzWIHIAweoPawLW1sRXiK43YfXisHJ1T2pjhmMwgTp00OOYoG5MoLLUl4HqQJHUA2UFlQjwIiAdwrBzyzZK+UIJ1hDYUwRxsAiDC0VVd7s4IarKYE2jaJVBziRnERDZCKKoysyADfMpEtLfKAKVwVpTuhDzZ66VgRtW3Zb+soiaQPBWFC7A9s2HCMgZAqyyaNoaXSNkn5mTCNLEsKkCrNEBVTj94nls23QQET7lBGs5XSgmU0SdTOlwCYFaDH+3QoqshXRByRciaVZKDW1nF5KiyimKpsHbKUuIqp5kIaraJ8PFNokgTRRAZVHV0TyGUmKClSKrp0M8XbRiCFSKrCZhAqgsHbYwCIa5QZAQVaOD0jGvfh6lwF4EqhelwF4cVq7uazz1WYcPeQyjhw2iqGoupL762o706ohdW5/F06tWYW7DbJTkV+Hci8dw4vAhikxJKMgox9kXTyI3pZIibDSlvwzUdAFUb+rvHdx7BFs27MSWp3Zidt1yHNx9BG++fhUb12zH4b0nkJM0HYf3HUdN2SLMrV+FG9fewUH6/MqZ13Hj6g2CeLEQMZfO3YRrr7+FF4+cQUHaLJw/fRmXL13Fivlb8NqFN7D72WMoypiNK6+8hS3r96Ewo0lY/rOPP0bDlBUUxauQHj8DV169imc37kU8XQzOtVygZa5i6Zz18HXJw47N+3D9yhWKsJm4ePo0VixZj2kV8/DysZNYsXA9ntu4Azkp1bQ/W/HchuewZM4KHNq1H80HDlCkjCAww3Fs3z4saVqCyIB0PL9zB3Y9s4VS4mBse3oj9j37HF0oEhAVlIrFTQsoU/EmUJkZrN4cVq7uSVNlFMYOHwDVUUMwaYKV0Ec1mOSMEL9YnD5+COtXPonnqMFtWruBGvUWrF22DhtWr0dL8zFkJpehurgRjmZRlOpm3RJRd23dj5NHTqD5hRO4cPoCDuw6is0E7xuvvY6rr72BGZVL8PabbyGM+o8H9x7Hm5evEZgbCaLrOPPSRQJzJooz5+Kl5nOUBs/HX999D9uf2SdAeuPKX+gisgtvvnYNjTWrcebkBYLsNSxbsBmvE8hZKTNx+uQ55KU1ISakChfOvIpkSqk/ePd9PLlkM9at2Ipzpy7g2Q27KeXNweXzr2L7lj0UdYvwwr4XUFU8my4kB+iicJb2eTvt/3O42PIiTp9oxqJZy3D80CHUVc/FC3v2CrDamERj61MbEeydiv3btlMfNoyWW4CmaU2IDqYLmEcc5tbNo0gdicykYoLUp9XeHFau7muC6mgojxgAtdHDoKNlC0MdJwHW1YuXY+HseZhRM4Mi4S4c2L4Dy+Ytxb5tu/HM+o0ESDNB+zSyUythbRSGDaueodQxldLLLAS65xAQW/DBX25gYdNavP7KG3h+z2FcPHMJL+xvFsBjEfcvN24g3KeMouUrePn4Gaxb/hz2bDtEEW+D0B+tKlxMcL+FmbVP4tLpVykq76LtbMPf//oRlhDYyxdsxN7tR/D6xTexZumzlLIvwI0336Z6P4WjB19EZmIDYkOn4pVzr6O8cA5t/1U8vXYH6qcsp4vDm1gyewNF0zxcOHUexw83o2nGciyb/xQBXoKP3n9LiKSXL5xHTfkcvH7hDDau3oiE0AKcOt6M3LQqvEhZhr1pFMEaiYM7dyImNBeHCeDwgDRs37QVM8rrkRabh+KscuSnT6Y+bwQBnQAzfT8yh5Wrh9JWH0ugDoSm8lgY6DhSf8pVeKwQ6hsPw0keBK8nwgNTEOafQpEhGjHB2fBwikNSdCE1zjyE++XAzyWVQKmjVwYq9VE9cnHsUDNFzjdpWib1PStRUTCPUuY6RPoXUzRegHL6XJI1W7hzW1W0kFLZJmTE1mNew3qKSKvpfR2yExtRkb8QCxqfRmZcPXKTZmH29KcoKq9Cbeky6v/RxSJhJvJSZmN+w9PIim9Aac4CSq2XE2ArkBJTT/3EqRRVZ6CBIjB7pJMeT9OCqwigWYgNrkQI9V+jg8pQmDkTb1+9RhlCPnyc81CUU4/IwBIUZs2Al1M64kInU9l1lMoW0r5PpmkpVG4J9VkTqBsQT/tYQOuxm2tJlOJTH9knnfr14UiPK6F+eA4s9APh5xaHMN8U6sP688jK1XPpaSpj/Jgh0FbXbLvzy25+mOixGyA+MDf0haVxIGyp/2VnHkYpXgScrKPhbBMLN1v2XDUJno7J8HFMRUBrVGV3fLOSa6lhFlGaW0TRs5hcgki/ycLNnpjASsQSOAzU+JAaJIWzu7nsEcwMpEezu74NyIyVc1wjvcqcQZ8zYhppuQakRjbQOnVIiahDcjh7rad+bh3BMR2xIdPpwlJLINaQp1BfsgoR/lV0cakkkMqEG06hXqUI9ihBIAEb6F6EIM8i6lsWURqfS33ZXAKPXp1zaN+y4CXcbEqHh32acMPJXbg7nARX20S42JCtE+BiFSN7nGPO7g5Hws6U+rMmYbBij3CM2V3hQFgY+guwyu4IewmPw5Q4rFzdkcEEZWiMpRRYQ1t4nmqq7yE0JNkNEB9qXH6wooZmaxbSCmsUwRoDV5s4YZCAl0My9ftYRGEDILIFWFlk9XFmAyHyCAgGa4ng6IAyMoO1Sg5WdgeXDWyQPX5Ji64XnE7OJCiZs2JnEqgzZZCS0wjUtCgCNaKBIK1HUtgMMgHLoI2YgfiwaRJYp9JFopqiOhskUUEXkzKCspSAJVg9b8Lq5yJ77hroWgA/1zz4UYT1c6J9ccyEt30mvOwJVgasQ6oArPAoxzapDVbn1mevThbRcGSDJcwiKUUOlwFrHCw8xrEwDGiFVTZAgt1kUuKwcnVHhloq0Bw3EgZaxjA1YCNqvIVoam7A7Cc0LjYax848HPYWkXC0YlE1jmCNF545skbr5cBGK2VSfzVbcJBnLkI8CwRQQ3wpqvpORqRvqTA4gY0qig2qFsxgjQtmI4/YCCQG7HQBNhYt0yLrkRHdGkUlkKZGNgpOCmOgNrSaRVf2uY5SVhmoMlhrERU4FdH+UxDl1wqrL0tTWXS9CWuQOxskUSg8d2X2c8kX7Euw+jrlwNuBgG2NrB52FF3t2MimlDZY2bNXJ0qHHS3j2mC1p8hqR7DaGjNYQwSLwJoZyJ63cli5ui2jiarQUhkLI21rgtOTQJWlvuYUUVnaJouqYQKosqgaSw2TQKWo6kYpoAeD1TFNeOzhx2ClPmughwzUMO9i4W5vpH8ZgVompL+yqMpgZa9TEBtYjXgBWDaggUXXGW3ApkaxKNsgmL1PITNIUyJmEpgySBND69ucQPDGE6yxITMQFzSdLEbWmtY0uJLSYDaaiSIrpcChXpPbYA0gWNmzVz8hqlKEJVCZb8KaRWlwJu0vG4KYRqDeCqsIqqNFnDCayc4sWoistqZsRBOLrqFtwAqw6t8cHKHEYeXqjlhknaiqAmMdWwK0NaoSqGZG/kI/i0VVNiKHgcrsSLD6O0YhwDmK0uAkYXSPj1MafF2zCNYsgjXnFlhZ+hvlz6IqS4Erb0ZWgjQyoAbTs9IJ4GkCsAxW+egqAis6uS31pahKsCYyWMkM0gSCNT6URdX6DmEVQG1NgxmkIZ4M1snCIAmWBjNYRWBvwppHaXB2K6xZt8DqaptCoLKBEjejqoN5rDBgxI5AFSKrMPzwJqyy4Yc3YeWRlavbMqbIqjNeEya6DrAwZg/p/ZFs6IEDZq7wNKL+FcHKRuLYsyF1FDGivSOwOF4PJX7mcKSIwm66+DqmI8E/Fb4uORSh8gjWPASz0UreDNLine truncated
+
+[image10]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAXYAAAFiCAYAAADiGfPoAAB7FElEQVR4XuydB1gV19aGc2tuv3/KTUwxsSb2rthFELtSFMUKduzYC1aw94rSRFBBUUBFEVEEsffeYktssbfEWKJ+//4WzvF4wIIVcH/Ps55zZs2ePYXDO2vW7PIBtLS0tLQylT6wdGhpaWlpZWxpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBrvbJOnDiBDh06IEuWLPjkk09QpkwZxMTE4N69e5ZF37muXbsmx2dtbY2HDx9artbSyhTSYNd6JY0YMQIffPBBqlaqVCmBfnrS1atXkTdvXhQpUkSDXSvTSoNd66UVGhpqgnixYsWwefNmHDt2DD169DD5K1WqhN9++02i999///2J7R88eJDCd//+/RQ+LtNY3lzG9gQ0zSj3LKUGdvPtUquDx0RLTUZ54zhSk3mdxvlZnotR5mn70dJKizTYtV5Kd+7cQenSpQXe/LQE1ahRo/DHP/5R1q9evRr//ve/5fuyZctkPSFYoEAB8a1bt058BNv//vc/8d29excnT56Eo6Oj6Sbx1Vdf4fDhw6Z9uLm5ib9hw4ZyYzHKtW3bVlIuqckS7KdOncIf/vAH2W7q1KmmOpo2bYpffvkF48aNk+U//elPst4c3tHR0abyxnXYvn27af2NGzfQrl0703pnZ2dUrFhRvk+ePNlULjIy0lTmP//5j1w7nr+W1stKg13rpXTw4EEBEaG4atUqy9USpWfLlk3K9OzZE35+fvLd1tZW1u/cudMEMwKaioqKkuXixYtj3759JuDyCYCAtbKykuXExEQp37FjR1Mdffv2Ffvb3/4my23atDEdi7kswX7mzBnTDejDDz/E2LFj8fXXX5vq5bFw33/+85/leBYvXiz1+Pv7y/r4+HhZDgkJkeWPP/5Y9nHz5k3897//FR/3xzrc3d1N9c6cOVO243FyuXXr1lKGuX8uFy1aVK6hltbLSINd66WUkJAgAGIku2nTJsvVknr57rvvpAyBxuibAP373/8uwDJy8/R98803EvG3atVKfN7e3ujSpYt8b968uUTfBOWWLVvER9BTBti5H0NNmjQRn4uLi8lnrmeBPSAgQMrwBsHl//u//8OFCxfElzt3bvEx/UTdvn0bt27dkjp4bEuWLJH1fHnMbfjymDeCf/7zn3LuhniDYDne6BiVcx9cvnLlitRz9OhR09PNgQMHTNtpaaVFGuxaL6VDhw4JfAivuLg4y9UCPiNinzZtmgDQSEMwHWNnZ4d//etfqFmzpvi2bdsmqZm//vWvcqNwcHAQf2pWvnx52YcBdnt7e9N+27dvL75GjRqZfOZ6Fth37NghZYz0C1/+GqkX4yYVFhZmqisiIkKOkzcr49gMsDPVwmWmiMzTN4ULFxY/I/5Lly7JU4Ll+RnGVI+W1stIg13rpcQI88svvxQAeXh4iI8Aq1evnsCSKQpC+h//+Ad2794t6728vEwgzpo1K/LkyWOK3F1dXSV6zZcvn+S2jdx6586dpT7mws+ePYuff/5Z9k0ZYGeO3ZCR7ngZsDP9QzEdw+WSJUs+FeyEP29qjMC5r9GjR8t6A+w+Pj6yzHM0fyHKfdNPsF++fFnAzqce4xz5yXOk8T2GltbLSINd66VE4BnRMSPvBQsWSPqFIKOPLzr56enpadpm48aNphw4jfBes2aN3AAMX4sWLaRsnz59ZJkpkIsXL4qPTScJxL1798qyAfYGDRqY9vEqYDfqfR7YCeoKFSrIsgHtDRs2yDLBzuPluRrvCJYuXSplmHoxfEzF8Hp9+umnsswUDMWbF1NCQUFBcoPT0noZabBrvbTOnz9vitoJR0bbRqsWw9gZyGgpQrB9//33pnVM0TBqNeBGW758uZQ9fvy4KU3BfRQsWFCgyWWjvncFdvr4Etg4Fp4X00Nc5svT06dPC/D54pU+5tl5/N9++63pPI2Xp0OGDJFlXjeW+fzzz2WZzUR1xK71stJg13plnTt3Dl27dpXWJART9erVpZWIAUOmUwzNmjVLyhByRkqFL0Ppq1y5sqkcRThOmjQJ2bNnl/VsTmjexpwvOeln80ZDvXv3Fh8Bn5quX78uECaYCWge+xdffCHbsKUPxTQKl6tVq2YCe9myZcXHljsUj437582GrXUYsbM5I8vwqcIQmzLmzJlT9sGnGsKb14TXwRCPgWW4LdfzKUZL61Wkwa71RsWUQnrrffo2xDbsQ4cOlTbp5m3bjRYvvBFoab0pabBrab0B8cnAaC3DF6wcR+ezzz6TZUbmT+tApaX1OqTBrqX1hsS27IUKFTLl1Wls/mh0sNLSelPSYNfSeoNip6MffvhB2v3TdEsXrbchDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUymNw72n3/+GQcOHMDhw4fTrZ06dcrysLW0tLQyrN442Ldt24YCBQogd+7c6c9y5ZaJkhMSEiwPW0tLSyvD6o2DfdeuXciZKydGBAzDpAUTMDFsfLqxyeGT4NzaGevWrbM8bK10pNu3b8t0cosXL053xtmRtLTSm94K2HPkzAGfxdOw4XwS1p5Zk25s06X1cGjuoMGeznXmzBm0b99eAP/rr7+mG7ty5QpGjx5tebhaWu9cGuwa7OleBHubNm0s3e9cvNFosGulR2mwa7Cne6VXsP/2228a7FrpUhrsGuzpXhrsWlppkwa7Bnu6lwa7llba9NbBvu5cIjZcSML682tlOelsQgrgvi3TYM8Y0mDX0kqb3irYN15ch4FTPNGqR0ssO7AU/cb1wZDpg5FwKh7L1fKKQ8uw+sRKLNkThVXHYwX6SecSsOLwMsQeXi43BQKZ62NUWa5fdSxWLUdi5dEYKRt/Mg5L90YhUdXJ5dXHWV+k1JtwarXaLhrR+xYj8XR8qmDnC7HAwEDEx8ebnYXWu9TLgD0kJAQ5c+bEZ599hhw5cqBjx46WRZ4QW7i4uLhYup8pDXat9Kp3AvZeo3ti373dGDV7BDoP6giPoV1QuqIVBk31RHXHarBvXBe2tSoL0BdsCkXdRnVQsWoFLNq2AF29uqKWc02UrmAlkHZs6oD6bvVQt3FtKd+oXSNUc6wq++HNoUptW9m+qr0d5q4NwbfZvoVr52aI/zEOmy9vELDzGO/evQtfX198//33+NOf/oStW7danorWO9LLgN1QlSpVTN8J+759++LYsWMIDw/HzZs3xT9//ny5mfMGcP/+fVP550mDXSu96q2DfXiAt3RYKlKyCKxrWCN6/2J0GdwJg6cNkvVNOjTG1qsb0bZ3Gwzz90KP4d0wOXwiVv4QgzmJs1FZbcPIfXrkVFXnVNRr7iTRNyG/4nA0qjtUk2h9xZFl0gmp04AO2HFzK5p3aga/5TNRzaGqKdfPiL1eCye0bNkSefLkwQcffCBGGBw8eBA7d+7E7t27sWfPHuzfv198HILg6NGjOH78OE6ePImffvoJp0+fxtmzZ3H+/HlcuHABly9flgjw+vXruHHjBm7duiUQeJrxKeFpdufOnWcab0hPs3v37qVqv//++1ONYHuaPXjw4Kn28OHDZ9qr6Ny5cy8N9kqVKsm1iIqKQpcuXZCUlIQKFSrI37B3795YsmQJpk+fLh2N6E/LsWqwa6VXvXWwS8Q+qid2/7YD265vxrqfE9F9mAcmhU9A/Ik4tOjmho2X1sFzUn94zRiCVt1bwCdqGrZe24SI7QtRupIV7JvYo55rPSzbvwQzlk5HdadqaNaxCVafjMP8DfMkOh8ROAwj1E2kx4ju2Hp9k9pnD8xc6gOXNg0lv2+AvX6LeqhXrx4++ugjE9hz5cqFfv36oUePHujWrRs8PDzQuXNndOrUSTrKtG3bFq1bt0arVq3QokULuLq6olmzZmjSpAkaNWqEhg0bwtnZWep1dHRE3bp1YW9vL+bg4CBGP83JyUmMZevXry/WoEEDMdZDY4qAxrpp3I9hTZs2FeP+mzdvLsbjcXNzE+Px0XisNB43jaCk8Vxo7dq1g7u7uxjPsUOHDmJMYfC8abwGNAKya9euYrw2NF6n7t27yzWj9ezZU6xXr15ihGifPn0kYjasf//+Yp6enmIDBgwQGzRoEAYOHCifgwcPlv3w+ryMCHbezLhvXq/x48fLcBKUv78/6tSpI995U65Ro4b5ps+VBrtWetU7AXtnFaEb+XLmwbt5d5Xomi9VrWtUknw44cwIu8+Y3ug/sR96j+mFoT6DUc2xmmw7ft5YTI+YqsDdDbtubYd7v7YCbkb6i7YuUJ+tVYQfLPWs+XEVKtpVQHB8EBq0cn4C7EzFbNq0SaJuAurvf/87smbNikuXLlmeyjuXEf0yQmYEbUTZBBejUkbwRqRP6PBJgcZekr/88oukHmh8iuDTxLVr18SuXr0qTxh80uB5X7x4UYxPHwQeB3KjMXLmkwlTIzReMw6gRuOTy48//ihPMSdOnJAnGhrTHoyOGRHTjhw5Ik89hw4dEuNTEAeJo/GpaN++fdi7d688JdH4xBQbGys3q5eRAXYvLy9MnDgRGzZskCcxasyYMXKz5fXi+Wqwa2UWvV2wX1iHKQsnYeSs4abWMPwcEzwSs+ICJHof5jcUFezKq4i9n9wIWLine truncated
+
+[image11]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABM8AAARlCAYAAACukOBGAACAAElEQVR4Xuy9B7hV13WuHSd5ch3Zia17rcQl13+K88dyXOLYlmPLRb13CVUkUEMIEB3Rkeiig+iiI3rvvffeO4gmeu+dM+56xz4TbRZ7IxDt7M33Ps98zll97rXGnGvOb4055l+YEEIIIYQQQgghhBAiJX8RXyGEEEIIIYQQQgghhEgg8UwIIYQQQgghhBBCiDRIPBNCCCGEEEIIIYQQIg0Sz4QQQgghhBBCCCGESIPEMyGEEEIIIYQQQggh0iDxTAghhBBCCCGEEEKINEg8E0IIIYQQQgghhBAiDRLPhBBCCCGEEEIIIYRIg8QzIYQQQgghhBBCCCHSIPFMCCGEEEIIIYQQQog0SDwTQgghhBBCCCGEECINEs+EuAhnz5jlKCnF09k0Kb49/B9fH98/1bazSfaXYhvr025PPm98e3xd8nJ826VcK36eVNeN/e/nih+Xajl+vaT1550jvu+lbIv/TfV/fDn5nPHz5yYhhBBCCCFE9iHxTIh05JgdP5Bjx3YrKcXSntTpeG5KXk61/lzafZFtsXPE07n1KfJz7pwp8nvB9ZL2uWBb8j7ptuWm+DWT973geqmOS943zfW+7Den+j++HD9Hym25/8eXk895wfmj8506mhOvRYQQQgghhBBZgMQzIdKBeLY/x45sz7GjO5SUvkhH0qT49vB/fH18/1TbLnbclx2bvC2+Pb4ueTm+Lb5Pqm2p9onvG///Upfj10ten25b/P/4cvzcqbalukb8nBecf7vEMyGEEEIIIbIViWdCXIQT+y3ROd6ppKSklD5RT0g8E0IIIYQQIjuReCZEOnIQz3ISHigpOstKSkpKIUk8E0IIIYQQInuReCZEOsKwTYlnSkpKX5IkngkhhBBCCJG9SDwTIh3yPFNSUrrEJPFMCCGEEEKI7EXimRDpkOeZkpLSJSaJZ0IIIYQQQmQvEs+ESIfEMyUlpUtMEs+EEEIIIYTIXiSeCZEOiWdKSkqXmCSeCSGEEEIIkb1IPBPiIpzYbxLPbmA6vjt6CEej57Dni+dgh83OHkgsu2CxL9rnpPlflk9wzIlon4Nf7HNyb2IfO56bonMe22UXXC+euD7nuZR9SX79KK/k8cuOYd/TkX3lHJKNZUOSeCaEEEIIIUT2IvFMiHRowoAbmhC8tq3bby0afWKfLd3mQhPrenToZ2MGT3ax7PRBs3lTllmJImVsxbz1dvpQ4pjypSrb8H7jfDv7LZu91t59s5gVzk1Vy9ew3RtQuOycoIZQhuCFsBYEtu1r91vfbkNs/+fsYC7K2ZGEKOYiGeuOJY4NQt6GKK9d2va0PZuOfnGu4wkB0MXAIOBFl58zabGNHDDBBTpSOD/58GXyd+ILYZB74PvkCm5nDnyxHL9/Stc3STwTQgghhBAie5F4JkQ6NGzzhiaEorWLt9hD9z9mcycvdVGJde++Ucwa1m7m4hMi1IAew+373/uBNW/Y1h/bkN6jfblRneYuPuHZNbDHCLvvrgds5viFtmDaUlsya5Xt2XjE5kxaZD079bcB3YfZwa2nXOBCdOvZaUC03zL7fNVuGzFgvK1bssUmjphhndv0sCmjZrvnG/v36NjPBvYc4fk9sSfHBa/5U5ZZwfxv2cblO23RjOXWu8tA69NlkO3//IQd3n7Gl7u172N7N+2zSmWq2aMPPm57Nh2wxTNXW6fW3WzCsOkuni2ZtSbK93Dr3qFvlI89djr6vRuWbY/y0N1mTljowtyaRZs9T3OnLPU8xe+h0vVLEs+EEEIIIYTIXiSeCXERNGzzxiWEMjzO8j39ki2YtsLFMtaVKVbeWjRs66IYAtKwvmPtyUeftVJFy9mhHcetecM2lu+pF619y09dcEPQGtpnjL30XH4XoXatP2CHo/3WR+f+nzvudE+0P995t/XpOtjGD5saXe9Fe+O1QvbgfY/YzAkLrGzxCtamWQe741f/YyXeLWP3/Pk+Wz5njbVq3N5KFStnBV9509o07eBeYORp0fSV5zzh7r/nQXvztXeiY+73/HRq3d3ee6ekVS73gQtolcpWs8ceftKFOfLfpF5Le+bxfDZ93Hx75YUC9nz025945Gmr8n4N2711X/S3ur1fspIVi84xacxUK/1eeav9QT17KV9+mzNxkZ2WgHbDksQzIYQQQgghsheJZ0JcBIlnNy5dqng2uNcoK1mkjL1fqrIN6D7CalerbzUq17HWTTu4eMZ+YwZNstv/4z8t/wsFrMDLr1vLxm1s7aLP7bmnXrAd6w5Yp1bdrH7NJlauRAUfErp1zT4r9EYRmzZmrpUrWTHa/xN7r3ApP1fZaJ9P2/exJx992lo1aW8fVKhlTz/+nB3YctLzE8SzZXPWuaj12dKt1rfrYKsV5avtx52sepS3kQNG277NB21g9+HWqM7HdmLP2SjvQ21Qr5H2wrMvW99Ph1iRQiVsanT9GePmW8WyVaPtw/28eNMd2XPUOrfuYX+68y4b2meYPfLA4+5pR/7i91Hp+iSJZ0IIIYQQQmQvEs+ESIeGbd7QhFC2bulWe+HZV1zock4kxLMOLbuee0yDeo6wmlU/cq8uvLVqVatnrZu0TwzjzPU8Q2B7u0DhRKyxXH1j3eIt7nW2c/2h6HyfutcXItzEETNt+7oDVrRQiYR4VqKifdK8k4tfDKdkqGW39n3tuadfsJpVPrKPG7T0oZOHtp2xnMNfiGcMu+TvphU7rVfHAVbnw4Z2cNux6Nge7s3W9ZNePuyzWb1WNnviYnv5+VetWf1W9uC9D/tQ0MrlPrSF01fYhOHTo99X14eSFn+39Lk4bd2jPNz9p/tcCKxfs7FNGjnTTu298D4qXZ8k8UwIIYQQQojsReKZEOmQeHZDE+LZ+mXb7Q+/v8uFq95detnMCXOtyFvF7eV8r1qvzr1sxMARLnzh/TV2yBS77Tu32bC+Y6xFo7ZWv1ZT9wTDGwvx7Df//VsX2Lq1722Deg6xuZOXWIGX37Bta/db8wZtrHHdFja490h7q0Bhq/NBA3vmiXw2bew8K/p2cR8KWrFMNc8XotqwfmPdY6xhnWY+bJLjT+xNzARKzLO3Cha2hdNX+vDPz5ZstS5te9iHlWpbuxado3O19mM/ad7Z+n06xAq88oa1b9HFr9e/21B75IHHXIwrWbSsTygwatBEK1uivK2cv87PhyhYqlhZ69Wlr4tzvToPsBLvlrbJo2b5PYvfR6XrkySeCSGEEEIIkb1IPBMiHRLPbmg6vjvH9m85YZ+262N1PmxgtarWs8G9R3nA/no1Gvsyccfw2po1YaHt3njExadd6w/ZwmkrfBbOk/uYtTPHA+03rdfSaler555pzeq3ttULNrlXF9fAW2zBtOW2d9Mxa1i7oRV45U2PZcZxU0fPskUzVvrwSfJFfLItq/b49Zi4oEGtprZpxS47tS8xacC2Nfv8vMzUOX7YNNuz8agtn7POxbrN0X51qzf0vByIrrt55S6Pp8YQT2boRFQbPWiSX2/K6Nm2dfVej82GMMaEALMmLvJhqUP6jLacI4nZOvk9vToNiOz0rN+z+H1Uuj5J4pkQQgghhBDZi8QzIdIh8eyGJ4ZJMlTTTuemI5aYZTMsn0gMy8yJ0ond5sMZmTGT4P2ITTw70sm9SceQTpqd2pc4F9dgf87F7JYli5axd94o4gIY52afMBkAeWKZ83EdO5VILIdr+frDudujv8d3J2K1cS7PB8ecTKz35ZOJ7f47o//5jZ6fQ4lznGKfQ4nze5y33PtwlGXyx/KxxO+I3z+l65ckngkhhBBCCJG9SDwTIh0Sz266hACFCIewpiGQSpeTJJ4JIYQQQgiRvUg8EyIdEs9uyoSAJi8upctNEs+EEEIIIYTIXiSeCZEOiWdKSkqXmCSeCSGEEEIIkb1IPBMiHRLPlJSULjFJPBNCCCGEECJ7kXgmRDoknikpKV1ikngmhBBCCCFE9iLxTIh0SDxTUlK6xCTxTAghhBBCiOxF4pkQ6ZB4pqSkdIlJ4pkQQgghhBDZi8QzIdIh8UxJKU8kZj89nsdnQJV4JoQQQgghRPYi8UyIdCCe7Ys6xdsTHWMlJaXrn45G6cRus5N77YJteSpF9cSpIxLPhBBCCCGEyEYkngmRhpyoH4wnCR1iJaXz0lG7cDms+7JtydtTLce3pTrvxY5LtRzfFj82+RrxY+P7xrelOzbVvun+j58nadvZE2Zd2/ewMsUr2IlDZ+x0/Nj4NeLnSZe/VMelOybVtvj/UTpzQuKZEEIIIYQQ2YjEMyGEEHmacePHWdu2bS0n52x8kxBCCCGEEEJccySeCSGEEEIIIYQQQgiRBolnQgghhBBCCCGEEEKkQeKZEEIIIYQQQgghhBBpkHgmhBAiTzNixAj76KOP7OxZxTwTQgghhBBCXH8kngkhhMjTFC9e3G699VY7ffp0fJMQQgghhBBCXHMkngkhhMjTVKxY0f7t3/7NTp06Fd8khBBCCCGEENcciWdCCCHyNJUqVbJ//dd/lXgmhBBCCCGEuCFIPBNCCJGnqVy5sv3Lv/yLxDMhhBBCCCHEDUHimRBCiDwNwzZ/+MMfSjwTQgghhBBC3BAkngkhhMjT4Hkm8UwIIYQQQghxo5B4JoQQIk+za9cuW7NmjeXk5MQ3CSGEEEIIIcQ1R+KZEEIIIYQQQgghhBBpkHgmhBBCCCGEEEIIIUQaJJ4JIYQQQgghhBBCCJEGiWdCCCHyNK1bt7Z77rlHEwYIIYQQQgghbggSz4QQQuRpqlSpYn//939vJ0+ejG8SQgghhBBCiGuOxDMhhBB5mpo1a9p3v/tdiWdCCCGEEEKIG4LEMyGEEHmaWrVq2T/+4z9KPBNCCCGEEELcECSeCSGEyNMgnv3DP/yDxDMhhBBCCCHEDUHimRBCiDxNtWrV7JZbbpF4JoQQQgghhLghSLine 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

@@ -0,0 +1,509 @@
+# Architecture Design Principles
+
+## Overview
+
+vLLM-omni is a multi-modality extension for vLLM that supports non-autoregressive structures and non-textual outputs. This document outlines the key software abstractions, APIs, and dependencies for the system, designed to maximize reuse of vLLM's proven architecture.
+
+## Architecture Principles
+
+1. **vLLM V1 Compatibility**: Built on vLLM's Engine V1 architecture with AsyncLLM and EngineCore patterns
+2. **Stage-based Processing**: Models are divided into multiple stages, each processed by different Engine Cores
+3. **Multiple Engine Core Support**: Each stage can use either AR Engine Core (reusing vLLM) or Diffusion Engine Core (new DiT support), or other new Cores
+4. **Worker Process Pattern**: Follows vLLM's multiprocess worker architecture for scalability
+5. **Extensibility**: Easy integration of new modalities, model architectures, and output formats
+
+## Key Data Flow
+
+API Server --> OmniLLM/AsyncOmniLLM (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
+ --> Scheduler (New one for DiT) --> Executor (New one for diffusers) --> Worker (New one for DiT)
+ --> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessoer (New one for final multimodal output)
+
+## Core Components (Ordered by Data Flow)
+
+### 1. Installation
+```bash
+pip install vllm
+pip install vllm-omni
+```
+
+### 2. Online inference launch (Entry Point)
+
+Keep consistency with the vllm main branch
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8000
+```
+
+#### Design Architecture
+```mermaid
+graph TD
+    A[vllm serve model --omni] --> B[vLLM-omni CLI Wrapper]
+    B --> C{Detect --omni flag}
+    C -->|Yes| D[Parse OmniConfig]
+    C -->|No| E[Forward to vLLM CLI]
+    D --> F[Initialize AsyncOmniLLM]
+    F --> G[Start omni Server]
+    G --> H[Multi-stage Processing]
+    E --> I[Standard vLLM Pipeline]
+
+    subgraph "vLLM-omni Components"
+        F
+        G
+        H
+    end
+
+    subgraph "vLLM Components"
+        I
+    end
+```
+
+#### omni Serve Command Implementation
+File: vllm_omni/entrypoints/cli/main.py
+```python
+import sys
+import argparse
+from typing import List, Optional
+from vllm_omni.entrypoints.omni 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 process with vLLM-omni
+        omni_args = [arg for arg in sys.argv[1:] if arg != "--omni"]
+        omni_serve = OmniServeCommand()
+        omni_serve.run(omni_args)
+    else:
+        # Forward to original vLLM CLI
+        from vllm.entrypoints.cli.main import main as vllm_main
+        vllm_main()
+
+if __name__ == "__main__":
+    main()
+```
+
+#### vLLM Plugin Integration
+File: vllm_omni/pyproject.toml
+```python
+[project.scripts]
+# Override vLLM's CLI entry point
+vllm = "vllm_omni.cli:main"
+
+# Add entry point for vLLM integration
+[project.entry-points."vllm.plugins"]
+omni = "vllm_omni.plugin:OmniPlugin"
+```
+
+### 3. Offline inference launch (Entry Point)
+Design an upper level class to incorporate multi Engines, each engine has a engine call.
+```python
+from vllm.entrypoints.llm import LLM
+from vllm.v1.engine.llm_engine import LLMEngine
+
+class OmniLLM(LLM):
+    """Extended LLM supporting multiple engines and stage-based processing"""
+
+    def __init__(self, stage_configs: List[StageConfig]):
+        super().__init__()
+        self.stage_configs = stage_configs
+        self.engine_list = []  # List of AsyncLLM instances for each stage
+        self.output_processor = MultimodalOutputProcessor()
+        self._initialize_stage_engines()
+
+    def _initialize_stage_engines(self) -> None:
+        """Initialize LLMEngine instances for each stage"""
+        for stage_config in self.stage_configs:
+            stage_llm = LLMEngine.from_vllm_config(
+                    vllm_config=stage_config.vllm_config,
+                    executor_class=self.executor_class,
+                    log_stats=self.log_stats
+                )
+            self.engine_list.append(stage_llm)
+
+    def generate(
+        self,
+        stage_args_list: List[stage_args],
+        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
+        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
+        priority: Optional[list[int]] = None,
+    ) -> list[RequestOutput]:
+        """Main generation interface - orchestrates multi-stage processing"""
+
+        # Process through each stage sequentially
+        for i, stage_config in enumerate(self.stage_configs):
+            stage_engine = self.engine_list[i]
+
+            # Prepare input for this stage
+            stage_args = stage_args_list[i]
+            prompt_str, engine_request, tokenization_kwargs = self._process_stage_inputs(stage_config, **stage_args)
+
+            # Add inputs to Engine
+            stage_engine.add_request(requesy_id, prompt_str, tokenization_kwargs)
+            # Run Engine
+            stage_output = stage_engine.step()
+
+            # Update stage input and output of each stage for later usage
+            self._update_stage_io(stage_output, stage_config)
+
+
+        # Process final output
+        final_output = self.output_processor.process_output(
+            self.stage_configs[-1].stage_output
+        )
+        return final_output
+
+    def _process_stage_inputs(stage_config, **stage_args) -> Any:
+        """Prepare input for specific stage"""
+        if stage_config.engine_type == "AR":
+            return self._process_ar_inputs(**stage_args)
+        elif stage_config.engine_type == "DiT":
+            return self._process_dit_inputs(**stage_args)
+        else
+            raise NotImplementedError
+
+    def _process_dit_inputs(**stage_args)-> Any:
+        image_latent = self.vae.encode(**stage_args)
+
+        return image_latent
+```
+
+### 4. Stage config for setting of different model stages
+```python
+@dataclass
+class StageConfig:
+    """Configuration for a processing stage"""
+    stage_id: int
+    engine_type: str  # "AR" or "DiT"
+    model_path: str
+    input_modalities: List[str]
+    output_modalities: List[str]
+    vllm_config: Optional[VllmConfig] = None  # For engine config of corresponding stage
+    executor_class: type[Executor]  # For execute class config of corresponding stage
+    dit_config: Optional[DiTConfig] = None  # For diffusion stages
+    cache_config: Optional[DiTCacheConfig] = None
+```
+For AR stage, the setting is:
+```python
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.executor.multiproc_executor import MultiprocExecutor
+ar_stage_config = StageConfig()
+ar_stage_config.vllm_config.scheduler_config.scheduler_cls = Scheduler # original vllm scheduler for AR
+ar_stage_config.executor_class = MultiprocExecutor # original vllm executor for AR
+```
+For DiT stage, the setting is:
+```python
+from vllm.v1.executor.multiproc_executor import MultiprocExecutor
+dit_stage_config = StageConfig()
+dit_stage_config.vllm_config.scheduler_config.scheduler_cls = DiffusionScheduler # New scheduler for DiT
+
+# For diffusion models without using diffusers
+dit_stage_config.executor_class = MultiprocExecutor
+
+# For diffusion models using diffusers
+dit_stage_config.executor_class = DiffusersPipelineExecutor
+```
+
+### 5. Online inference main class
+
+Similar to OmniLLM in offline inference, add some asynchronous processing, referring to AsyncLLM
+```python
+from vllm.v1.engine.async_llm import AsyncLLM
+
+
+class AsyncOmniLLM(AsyncLLM):
+    """Extended AsyncLLM supporting multiple engines and stage-based processing"""
+
+    def __init__(self, stage_configs: List[StageConfig]):
+        super().__init__()
+```
+
+### 6. Engine Core
+No need to change. The specific executor, scheduler and output type can be transferred into it with new configs.
+```python
+# class: EngineCore
+# EngineCore.step (simplified)
+class EngineCore:
+    def step(self):
+        scheduler_output = scheduler.schedule()
+        model_output = executor.execute_model(scheduler_output)
+        engine_outputs = scheduler.update_from_output(
+            scheduler_output, model_output
+        )
+        return engine_outputs
+```
+
+### 7. Scheduler
+Create a new Diffusion Scheduler for DiT, which is inherited from original vllm AR scheduler. At first, no complicated strategy.
+Just process the request first in first out. Then create a class of DiT Cache Manager for optimization in the future.
+```python
+from vllm.v1.core.sched.scheduler import Scheduler
+class DiffusionScheduler(Scheduler):
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        kv_cache_config: KVCacheConfig,
+        dit_cache_config: DiTCacheConfig,
+        structured_output_manager: StructuredOutputManager,
+        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
+        include_finished_set: bool = False,
+        log_stats: bool = False,
+    ) -> None:
+
+    self.dit_cache_manager = DiTCacheManager(dit_cache_config)
+```
+For DiT Cache Manager (Can refer to xDiT):
+```python
+class DiTCacheManager:
+    """Manages DiT-specific caching"""
+
+    def __init__(self, config: DiTCacheConfig):
+        self.cache_tensors: Dict[str, torch.Tensor] = {}
+        self.cache_groups: List[DiTCacheTensor] = config.dit_cache_tensors
+
+    def allocate_cache(self, request_id: str, size: int) -> torch.Tensor
+    def get_cache(self, request_id: str) -> Optional[torch.Tensor]
+    def release_cache(self, request_id: str) -> None
+    def clear_expired_cache(self) -> None
+```
+
+### 8. Executor for DiT without diffusers
+
+No need to change. Just import original AR executor.
+
+### 9. Diffusers Pipeline Executor (no worker)
+
+A single-process executor that directly runs the Diffusers pipeline without spawning workers or using RPC. Interfaces remain identical to `Executor` so the EngineCore loop is unchanged.
+
+#### Function map（Pipeline Executor）
+
+##### 1) Inherited and overridden
+```python
+from concurrent.futures import Future
+from typing import Optional, Union, Callable, Any
+import torch
+import torch.nn as nn
+from vllm.v1.executor.abstract import Executor
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.tasks import SupportedTask
+from diffusers import DiffusionPipeline
+
+class DiffusersPipelineExecutor(Executor):
+    supports_pp: bool = False  # Single-process, no TP/PP/DP
+
+    def _init_executor(self) -> None:
+        # Called by ExecutorBase.__init__
+        self._failure_callback: Optional[Callable[[], None]] = None
+        self._device = self._resolve_device()
+        self._dtype = self._resolve_dtype()
+        self._pipeline = self._build_pipeline(device=self._device, dtype=self._dtype)
+        self._profiler = None
+        self._is_failed = False
+        self.is_sleeping = False
+        self.sleeping_tags: set[str] = set()
+
+    # major functions to build/run diffusers pipeline
+    def _build_pipeline(self, device:torch.device, dtype:torch.dtype)->DiffusionPipeline:
+        model_name = "Qwen/Qwen-Image"
+
+        self.pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=dtype)
+        self.pipe = pipe.to(device)
+
+    def _run_pipeline(self, scheduler_output) -> ModelRunnerOutput: ...
+        positive_magic = {
+            "en": ", Ultra HD, 4K, cinematic composition.", # for english prompt
+            "zh": ", 超清，4K，电影级构图." # for chinese prompt
+        }
+
+        # Generate image
+        prompt_embeds = self._get_and_process_prompt_embeds(scheduler_output, positive_magic)
+        negative_prompt_embeds = self.pipe.embed_prompt(" ")
+
+        width, height = aspect_ratios["16:9"]
+
+        image = pipe(
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            width=width,
+            height=height,
+            num_inference_steps=50,
+            true_cfg_scale=4.0,
+            generator=torch.Generator(device="cuda").manual_seed(42)
+        ).images[0]
+
+        output = self.wrap_image_as_ModelRunnerOutput(image)
+        return output
+
+    # ---- Internal helpers (implementation-specific, not public API) ----
+    def _resolve_device(self): ...
+    def _resolve_dtype(self): ...
+    def _get_model(self) -> nn.Module: ...
+    def _get_and_process_prompt_embeds(self, scheduler_output, positive_magic):
+        ...
+        #append the positive_magic to prompt and embed them to prompt embed tensors
+
+    # Functions related to workers should either raise NotImplementedError
+    # or return empty defaults in this no-worker executor.
+
+    def collective_rpc(self, method, timeout=None, args=(), kwargs=None) -> list[Any]:
+        # No workers in pipeline executor
+        raise NotImplementedError("No workers in DiffusersPipelineExecutor")
+
+    def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
+        return  # no-op (pipeline already built in _init_executor)
+
+    def register_failure_callback(self, callback):
+        self._failure_callback = callback
+
+    def determine_available_memory(self) -> list[int]:  # bytes
+        # Single device; return [available_bytes]. If CPU-only, return [0].
+        return [self._determine_available_bytes(self._device)]
+
+```
+
+### 10. Worker
+
+#### Inheritance strategy
+Prefer reusing the mature GPU Worker end-to-end. Worker class is selected by configuration (vllm_config). Do not add a new executor-specific worker binding. If customization is needed, override only `init_device` to construct the `DiffusionModelRunner`; all other behaviors (device init details, profiling, sleep/wake, PP/TP comms, execute path) remain from `vllm/v1/worker/gpu_worker.py::Worker`.
+
+
+Inherited and overridden
+```python
+# Optional: only if you need to plug a custom DiffusionModelRunner.
+from vllm.v1.worker.gpu_worker import Worker as GPUWorker
+from vllm.v1.worker.diffusion_model_runner import DiffusionModelRunner
+
+class Worker(GPUWorker):
+    def init_device(self) -> None:
+        #those related to device check and init
+        ...
+
+        self.model_runner = DiffusionModelRunner(self.vllm_config, ...)
+```
+
+### 11. Model Runner
+
+### Function map (Model Runner)
+#### 1) Inherited and overridden
+Those parts relied to the KV Cache will be omitted if we do not register the model to the vllm config. The engine core will view it as do not require KV Cache, and handle it properly
+
+Reuse `vllm/v1/outputs.py::ModelRunnerOutput`：
+- DiffusionModelRunner: Use the `pooler_output=[Tensor,...]` to return multi modal tensors
+- ARModelRunner: Use the `pooler_output=[Tensor,...]` to return hidden states.
+```python
+from typing import Optional, Union
+import torch
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+from vllm.v1.outputs import ModelRunnerOutput
+
+
+class DiffusionModelRunner(GPUModelRunner):
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[ModelRunnerOutput, IntermediateTensors]:
+        ...
+        return ModelRunnerOutput(
+            req_ids=[...],
+            req_id_to_index={...},
+            sampled_token_ids=[],
+            spec_token_ids=None,
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[Tensor,...],         # Return Hidden states
+            kv_connector_output=None,
+            num_nans_in_logits=None,
+        )# return multi modal tensors via pooler_output=[Tensor,...]
+
+
+class ARModelRunner(GPUModelRunner):
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[ModelRunnerOutput, IntermediateTensors]:
+        ...
+        return ModelRunnerOutput(
+            req_ids=[...],
+            req_id_to_index={...},
+            sampled_token_ids=[],
+            spec_token_ids=None,
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[Tensor,...],         # Return Hidden states
+            kv_connector_output=None,
+            num_nans_in_logits=None,
+        )# return hidden states via pooler_output=[Tensor,...]
+```
+
+### 12. RequestState Role and Usage
+```bash
+# RequestState serves as the per-request state tracker in OutputProcessor:
+# - Maintains request-specific state (tokens, logprobs, detokenizer, etc.)
+# - Converts EngineCoreOutput → RequestOutput/PoolingRequestOutput
+# - Manages request lifecycle from registration to completion
+AsyncLLM.add_request()
+↓
+OutputProcessor.add_request()
+↓
+RequestState.from_new_request() → 创建请求状态
+
+AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() →  创建一个Background loop 持续从EngineCore获取输出
+↓
+OutputProcessor.process_outputs() → 更新状态并处理输出
+↓
+RequestState.make_request_output() → 转换为最终输出,格式为RequestOutput或者PoolingRequestOutput
+↓
+RequestOutputCollector.put() → 推送到队列（AsyncLLM）
+```
+
+Need to add implementation for an existing method
+```python
+class RequestState:
+    def __init__(self, request_id: str, parent_req: Optional[ParentRequest],
+                 request_index: int, lora_name: Optional[str],
+                 output_kind: RequestOutputKind, prompt: Optional[str],
+                 prompt_token_ids: list[int], logprobs_processor: Optional[LogprobsProcessor],
+                 detokenizer: Optional[IncrementalDetokenizer],
+                 max_tokens_param: Optional[int], arrival_time: float,
+                 queue: Optional[RequestOutputCollector], log_stats: bool):
+
+    def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput:
+        """Create PoolingOutput for multimodal/pooling requests"""
+        return PoolingOutput(data=pooling_output)
+```
+
+
+### 13. OutputProcessor
+
+For hidden state output, original OutputProcessor already support.
+
+We only need to add new one for final multimodal output.
+```python
+class MultimodalOutputProcessor(OutputProcessor):
+    """Handles multimodal output processing"""
+
+    def __init__(self):
+        self.output_handlers: Dict[str, OutputProcessor] = {}
+
+    def process_outputs(self, engine_core_outputs: list[EngineCoreOutput], ...):
+        for engine_core_output in engine_core_outputs:
+            # Option 1: Use output_type field (if EngineCoreOutput is extended)
+            if engine_core_output.output_type == "image":
+                self._process_image_output(engine_core_output)
+            elif engine_core_output.output_type == "text+image":
+                self._process_text_image_output(engine_core_output)
+            elif engine_core_output.output_type == "latents":
+                self._process_latents_output(engine_core_output)
+            elif engine_core_output.output_type == "text":
+                self._process_text_output(engine_core_output)
+            else:
+                # Fallback: use existing pooling_output logic
+                if engine_core_output.pooling_output is not None:
+                    self._process_pooling_output(engine_core_output)
+                else:
+                    self._process_text_output(engine_core_output)
+```

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/mrs_design.md

@@ -0,0 +1,188 @@
+## Multi-Request Streaming (MRS) on a Single Machine
+
+### 1. Background & Scope
+- All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved.
+- Current alignment with vllm-omni: `OmniLLM` supports multiple stages (`OmniStage`). GPU runners already expose streamable steps (prefill/decoding/diffusion), but the entry layer still collects lists and lacks intra-stage streaming and window scheduling.
+- Goal: implement multi-stage, multi-request streaming (MRS) locally. Each stage outputs segments; downstream stages stitch and trigger compute based on configured windows. Shared memory and zero-copy strategies reduce data movement overhead.
+
+### 2. Key Constraints
+- Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`).
+- Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version.
+- Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0).
+
+### 3. Architecture Overview
+- Processes & IPC queues
+  - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue.
+  - Inter-stage connection via IPC: copy-based `multiprocessing.Queue` passing dict payloads; use shared memory for large objects.
+  - Each link is SPSC (single-producer/single-consumer): the upstream is the orchestrator and the downstream is a single stage process; queues are unbounded (maxsize=0) on the orchestrator side.
+- Device visibility
+  - Each stage sets `CUDA_VISIBLE_DEVICES` or calls `torch.cuda.set_device` to bind to GPU sets.
+  - A stage may use multiple GPUs internally (TP/PP/DP) but presents as a single stage unit.
+- Simplified IPC: copy-based queues/pipes for data transfer; zero-copy is future work.
+- Pipeline progression: when a stage finishes a request, it enqueues outputs to the downstream stage; if downstream is idle, it starts immediately.
+- Scheduling
+  - A downstream stage triggers only after the upstream completes the request.
+  - Windowed segmentation/stitched triggering is not implemented; intra-stage streaming is not provided.
+
+### 7. IPC Implementation (simplified: copy-based)
+- Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data).
+- Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version.
+- Backpressure: queues are unbounded; pressure manifests as compute-rate differences. Optional SHM reduces large-object transfer cost; RX/decoding overhead is recorded for observability.
+
+### 8. Scheduling & Cancellation (simplified)
+- Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency.
+- Cancellation/timeout: explicit cancellation/timeouts are not provided; graceful shutdown uses a `None` sentinel sent to each stage input queue.
+
+#### Short sequence example (req0/req1, stage A→B)
+1) t0: stage A handles req0
+2) t1: req0 completes on A → enters B; A immediately starts req1
+3) t2: B handles req0 while A handles req1 (parallel across stages)
+
+### 9. Integration Points (by file)
+- `vllm_omni/entrypoints/omni_llm.py` (Orchestrator)
+  - Class `OmniLLM` orchestrates multi-process stages; constructs `OmniStage` instances in parallel and spawns per-stage workers.
+  - Spawns stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), creates control/data channels, builds simple full-trigger flow.
+  - Stats/logging are disabled by default; per-stage and orchestrator stats are only written when explicitly enabled.
+  - Manages process lifecycle: start/wait for readiness, graceful shutdown; forwards results between stages using copy-based IPC and optional SHM.
+  - Stage readiness: each stage emits `{"type": "stage_ready"}` after initialization; the orchestrator waits for all stages or times out and logs diagnostic suggestions.
+
+### 9.1 Process Device Visibility
+- Device binding: set `CUDA_VISIBLE_DEVICES` before process start, or call `torch.cuda.set_device` early in init; configs can be single/multi-GPU.
+- Cross-device transfer (simplified): via CPU copies; zero-copy paths are out of scope for now.
+
+### 10. Worker Template (pseudo-code)
+```python
+def stage_worker(input_q, output_q, runtime, shm_threshold):
+    # Device binding
+    set_stage_gpu_devices(runtime.devices)
+    engine = OmniStageLLM(model, **engine_args)
+    batch_seq = 0
+    while True:
+        first = input_q.get()
+        if first is None:
+            break
+        # Batch intake (up to runtime.max_batch_size)
+        batch = [first]
+        while len(batch) < int(runtime.max_batch_size or 1) and not input_q.empty():
+            nxt = input_q.get()
+            if nxt is None:
+                input_q.put(None)
+                break
+            batch.append(nxt)
+        # Decode IPC payload and sampling params
+        request_ids = []
+        engine_inputs = []
+        rx_bytes = {}
+        rx_decode_ms = {}
+        for t in batch:
+            ein, rx = maybe_load_from_ipc_with_metrics(t, "engine_inputs", "engine_inputs_shm")
+            request_ids.append(t["request_id"])
+            engine_inputs.extend(ein if isinstance(ein, list) else [ein])
+            rx_bytes[t["request_id"]] = rx["rx_transfer_bytes"]
+            rx_decode_ms[t["request_id"]] = rx["rx_decode_time_ms"]
+        # Generate and dispatch (grouped by request_id)
+        batch_seq += 1
+        outputs = list(engine.generate(engine_inputs, batch[0]["sampling_params"], use_tqdm=False))
+        grouped = group_by_request_id(outputs, request_ids)
+        for rid in request_ids:
+            r_out = grouped.get(rid, [])
+            use_shm, payload = maybe_dump_to_shm(r_out, shm_threshold)
+            msg = {
+                "request_id": rid,
+                "metrics": {
+                    "batch_id": batch_seq,
+                    "rx_transfer_bytes": int(rx_bytes.get(rid, 0)),
+                    "rx_decode_time_ms": float(rx_decode_ms.get(rid, 0.0)),
+                    # omitted: generation latency and token stats
+                },
+            }
+            if use_shm:
+                msg["engine_outputs_shm"] = payload
+            else:
+                msg["engine_outputs"] = payload
+            output_q.put(msg)
+```
+
+### 12. Observability
+- Metrics (as implemented):
+  - per-request (emitted by stages): `num_tokens_out`, `stage_gen_time_ms`, `rx_transfer_bytes`, `rx_decode_time_ms`, `rx_in_flight_time_ms`, `batch_id`
+  - orchestrator aggregates: E2E latency, tokens/s (written only when stats are enabled)
+  - optional per-stage JSONL: `{log_file}.stage{stage_id}.stats.jsonl`
+- Logs/tracing:
+  - optional per-stage log files: `{log_file}.stage{stage_id}.log`
+  - the orchestrator can log readiness, forward size/time, and summary information
+
+### 15. Risks & Mitigations
+— End-to-end latency: full-trigger sacrifices interactivity; start simple/correct, add windowed streaming later.
+— CPU copy overhead: copy-based IPC for maintainability; add SHM/zero-copy later as an optimization path.
+
+### 16. Refactor: Sink LLM init and process into Stage (diagram)
+
+Objective: without changing external call patterns, encapsulate per-stage LLM init, process creation, and worker logic inside `Stage`, so `PipelinedOmniLLM` focuses on orchestration (seeding, polling, forwarding, collecting) while keeping multi-process, shared memory, and device mapping capabilities.
+
+Key changes (vs. current):
+- Before: Orchestrator created processes directly, passing `_stage_worker`; device setup, LLM init, batching, and SHM were implemented there.
+- After: enhanced `OmniStage` owns its subprocess (`init_stage_worker`), with `_stage_worker` as the entry; device setup, LLM init, batching, and SHM live in `OmniStage`.
+
+Classes & responsibilities
+- OmniLLM (orchestrator)
+  - Build `OmniStage` list in parallel (preserve `process_engine_inputs` wiring)
+  - Connect adjacent `Stage` input/output queues; seed requests to stage 0
+  - Poll stage outputs, decode results, call `process_engine_inputs`, then encode and forward to next stage
+  - Termination & cleanup: distribute shutdown signals; join/terminate subprocesses
+- Stage (stage unit)
+  - Members: `stage_config` (with `engine_args` and `runtime`), `in_q/out_q`, subprocess handle, stats
+  - API:
+    - `init_stage_worker(...)`: spawn subprocess
+    - `stop_stage_worker()`: graceful exit
+    - `submit(ipc_payload)`: submit to input queue (may use shared memory)
+    - `try_collect() -> Optional[payload]`: non-blocking get from output queue
+    - `process_engine_inputs(stages, prompts)`: reuse/delegate existing logic
+  - Subprocess entry `worker_main()`:
+    - `set_stage_gpu_devices` (`CUDA_VISIBLE_DEVICES` & logical index mapping)
+    - build `OmniStageLLM(model, **engine_args)`
+    - while-loop:
+      - take first task + up to `runtime.max_batch_size-1` non-blocking; support `max_batch_size>=1`
+      - `maybe_load_from_ipc_with_metrics` per `engine_inputs` in batch
+      - `stage_engine.generate(batched_inputs, sampling_params)` (window `-1`)
+      - write per-request aggregated outputs to output queue after `maybe_dump_to_shm`
+      - log start/finish; report `error` per request on exceptions
+
+Data flow (sketch)
+```
+OmniLLM
+  ├─ Stage[0].start_process(in0, out0)
+  ├─ Stage[1].start_process(in1, out1)
+  ├─ ...
+  └─ Orchestrator loop:
+       - stage0.submit(encode(engine_inputs))
+       - for each stage:
+           res = stage.try_collect()
+           outputs = decode(res)
+           stage.set_engine_outputs(outputs)
+           if !final: next_inputs = next_stage.process_engine_inputs(stages, original_prompt)
+                    next_stage.submit(encode(next_inputs))
+```
+
+Config (preserve & extend)
+- `runtime.devices`:
+  - Comma-separated visibility list (e.g., "2,5,7"): set visible set; default logical index 0 (first GPU)
+  - Number/string: treat as logical index (0/1-based per implementation) mapping to physical GPUs
+- `runtime.max_batch_size`: max requests taken per intake (default 1)
+- `shm_threshold_bytes`: objects larger than this go via shared memory (default 64KB; injected into stages)
+- Stage link and inputs:
+  - `engine_input_source`: list of upstream stage_id(s) from which the downstream stage takes inputs (typically a single source)
+  - `custom_process_input_func`: function to construct downstream engine_inputs from upstream outputs
+- Output flags:
+  - `final_output`, `final_output_type`: marks that a stage produces the final external output and its type
+
+Kept & strengthened
+- Multi-process: each stage in its own process; parallel across stages; serial/batching inside a stage controlled by `max_batch_size`
+- Shared memory: large objects via SHM with strict `memoryview -> close -> unlink`
+- Device mapping: unified `set_stage_gpu_devices`, support comma lists and logical indices
+
+Migration (execution order)
+1) Enhanced `OmniStage` with `init_stage_worker/stop_stage_worker/submit/try_collect/_stage_worker`.
+2) `OmniLLM` holds `OmniStage` instances, replacing the older pipelined class; uses Stage APIs.
+3) Reuse `entrypoints/stage_utils`: device/SHM/encode-decode utilities live here.
+4) Regression: single/multi-request, various `max_batch_size`, SHM threshold, device mapping, multi-stage parallelism, error paths, and shutdown flows.

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

@@ -0,0 +1,97 @@
+# 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
+
+Use Docker to keep consistent basic environment (Optional, Recommended)
+```bash
+docker run --gpus all --ipc=host --network=host -v $source_dir:$container_dir --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
+```
+
+### Installation of vLLM (for users)
+
+Now we build it based on vLLM v0.11.0. Please install it with command below.
+```bash
+uv pip install vllm==0.11.0 --torch-backend=auto
+```
+
+### Installation of vLLM (for developers)
+
+If you want to check or debug with source code of vLLM, install stable release version of vllm with 0.11.0 from source with pre-built wheel file.
+
+```bash
+git clone https://github.com/vllm-project/vllm.git
+cd vllm
+git checkout v0.11.0
+```
+Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set VLLM_PRECOMPILED_WHEEL_LOCATION as your local absolute path of whl file.
+```bash
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.11.0/vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl
+```
+Install vllm with command below (If you have no existing PyTorch).
+```bash
+uv pip install --editable .
+```
+
+Install vllm with command below (If you already have PyTorch).
+```bash
+python use_existing_torch.py
+uv pip install -r requirements/build.txt
+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 successful.
+```bash
+python -c "import vllm._C"
+```
+
+### Installation of vLLM-omni
+Install additional requirements for 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
+
+## 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

@@ -0,0 +1,243 @@
+import argparse
+import os
+import os as _os_env_toggle
+import random
+
+import numpy as np
+import soundfile as sf
+import torch
+from utils import make_omni_prompt
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+_os_env_toggle.environ["VLLM_USE_V1"] = "1"
+
+SEED = 42
+# Set all random seeds
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+
+# Make PyTorch deterministic
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+# Set environment variables for deterministic behavior
+os.environ["PYTHONHASHSEED"] = str(SEED)
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        required=True,
+        help="Path to merged model directory (will be created if downloading).",
+    )
+    parser.add_argument("--thinker-model", type=str, default=None)
+    parser.add_argument("--talker-model", type=str, default=None)
+    parser.add_argument("--code2wav-model", type=str, default=None)
+    parser.add_argument(
+        "--hf-hub-id",
+        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(
+        "--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("--max-model-len", type=int, default=32768)
+    parser.add_argument(
+        "--init-sleep-seconds",
+        type=int,
+        default=20,
+        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
+    )
+
+    parser.add_argument("--thinker-only", action="store_true")
+    parser.add_argument("--text-only", action="store_true")
+    parser.add_argument("--do-wave", action="store_true")
+    parser.add_argument(
+        "--prompt_type",
+        choices=[
+            "text",
+            "audio",
+            "audio-long",
+            "audio-long-chunks",
+            "audio-long-expand-chunks",
+            "image",
+            "video",
+            "video-frames",
+            "audio-in-video",
+            "audio-in-video-v2",
+            "audio-multi-round",
+            "badcase-vl",
+            "badcase-text",
+            "badcase-image-early-stop",
+            "badcase-two-audios",
+            "badcase-two-videos",
+            "badcase-multi-round",
+            "badcase-voice-type",
+            "badcase-voice-type-v2",
+            "badcase-audio-tower-1",
+            "badcase-audio-only",
+        ],
+        default="text",
+    )
+    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).",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="outputs",
+        help="Output directory to save text and wav files together.",
+    )
+    parser.add_argument(
+        "--thinker-hidden-states-dir",
+        default="thinker_hidden_states",
+        help="Path to thinker hidden states directory.",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    model_name = args.model
+    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, 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}")
+    except Exception as e:
+        print(f"[Error] Failed to load prompts: {e}")
+        raise
+
+    if args.prompts is None:
+        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
+    omni_llm = OmniLLM(
+        model=model_name,
+        log_stats=args.enable_stats,
+        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+    )
+    thinker_sampling_params = SamplingParams(
+        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 = SamplingParams(
+        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 = SamplingParams(
+        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,
+    ]
+
+    prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
+    omni_outputs = omni_llm.generate(prompt, sampling_params_list)
+
+    # Determine output directory: prefer --output-dir; fallback to --output-wav
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    os.makedirs(output_dir, exist_ok=True)
+    for stage_outputs in omni_outputs:
+        if stage_outputs.final_output_type == "text":
+            for output in stage_outputs.request_output:
+                request_id = int(output.request_id)
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = args.prompts[request_id]
+                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+        elif stage_outputs.final_output_type == "audio":
+            for output in stage_outputs.request_output:
+                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)
+                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/qwen2_5_omni/extract_prompts.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+import argparse
+from typing import Optional
+
+
+def extract_prompt(line: str) -> Optional[str]:
+    # 提取第一个 '|' 与第二个 '|' 之间的内容
+    i = line.find("|")
+    if i == -1:
+        return None
+    j = line.find("|", i + 1)
+    if j == -1:
+        return None
+    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)",
+    )
+    args = parser.parse_args()
+
+    prompts = []
+    with open(args.input, encoding="utf-8", errors="ignore") as f:
+        for line in f:
+            if len(prompts) >= args.topk:
+                break
+            p = extract_prompt(line.rstrip("\n"))
+            if p:
+                prompts.append(p)
+
+    with open(args.output, "w", encoding="utf-8") as f:
+        for p in prompts:
+            f.write(p + "\n")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/qwen2_5_omni/processing_omni.py

@@ -0,0 +1,367 @@
+from __future__ import annotations
+
+import base64
+import logging
+import math
+import os
+import time
+import warnings
+from functools import lru_cache
+from io import BytesIO
+
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+from torchvision import io, transforms
+from torchvision.transforms import InterpolationMode
+
+logger = logging.getLogger(__name__)
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+VIDEO_MIN_PIXELS = 128 * 28 * 28
+VIDEO_MAX_PIXELS = 768 * 28 * 28
+VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
+FRAME_FACTOR = 2
+FPS = 2.0
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+
+temporal_patch_size = 2
+spatial_patch_size = 14
+spatial_merge_size = 2
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is
+    divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is
+    divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = IMAGE_FACTOR,
+    min_pixels: int = MIN_PIXELS,
+    max_pixels: int = MAX_PIXELS,
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            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))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        image_obj = Image.open(requests.get(image, stream=True).raw)
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        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:
+        resized_height, resized_width = smart_resize(
+            ele["resized_height"],
+            ele["resized_width"],
+            factor=size_factor,
+        )
+    else:
+        width, height = image.size
+        min_pixels = ele.get("min_pixels", MIN_PIXELS)
+        max_pixels = ele.get("max_pixels", MAX_PIXELS)
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=size_factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+    image = image.resize((resized_width, resized_height))
+
+    return image
+
+
+def smart_nframes(
+    ele: dict,
+    total_frames: int,
+    video_fps: int | float,
+) -> int:
+    """calculate the number of frames for video used for model inputs.
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+            support either `fps` or `nframes`:
+                - nframes: the number of frames to extract for model inputs.
+                - fps: the fps to extract frames for model inputs.
+                    - min_frames: the minimum number of frames of the video,
+                      only used when fps is provided.
+                    - max_frames: the maximum number of frames of the video,
+                      only used when fps is provided.
+        total_frames (int): the original total number of frames of the video.
+        video_fps (int | float): the original fps of the video.
+
+    Raises:
+        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
+
+    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`"
+    if "nframes" in ele:
+        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
+    else:
+        fps = ele.get("fps", FPS)
+        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
+        max_frames = floor_by_factor(
+            ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
+            FRAME_FACTOR,
+        )
+        nframes = total_frames / video_fps * fps
+        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}], but got {nframes}.")
+    return nframes
+
+
+def _read_video_torchvision(
+    ele: dict,
+) -> torch.Tensor:
+    """read video using torchvision.io.read_video
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    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.")
+        if "file://" in video_path:
+            video_path = video_path[7:]
+    st = time.time()
+    video, audio, info = io.read_video(
+        video_path,
+        start_pts=ele.get("video_start", 0.0),
+        end_pts=ele.get("video_end", None),
+        pts_unit="sec",
+        output_format="TCHW",
+    )
+    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=}, 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()
+    video = video[idx]
+    return video, total_duration, nframes
+
+
+def is_decord_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("decord") is not None
+
+
+def _read_video_decord(
+    ele: dict,
+) -> torch.Tensor:
+    """read video using decord.VideoReader
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    import decord
+
+    video_path = ele["video"]
+    st = time.time()
+    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.")
+    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=}, 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()
+    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
+    return video, total_duration, nframes
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if FORCE_QWENVL_VIDEO_READER is not None:
+        video_reader_backend = FORCE_QWENVL_VIDEO_READER
+    elif is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    # print(f"qwen-vl-utils using {video_reader_backend} to read video.",
+    #       file=sys.stderr)
+    return video_reader_backend
+
+
+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)
+        frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes
+        grid_timestamps = frame_timestamps[::FRAME_FACTOR]
+        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
+        nframes, _, height, width = video.shape
+        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
+        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
+        max_pixels = max(
+            min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
+            int(min_pixels * 1.05),
+        )
+        max_pixels = ele.get("max_pixels", max_pixels)
+        # min_pixels = (factor ** 2) * 52
+        # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
+        if "resized_height" in ele and "resized_width" in ele:
+            resized_height, resized_width = smart_resize(
+                ele["resized_height"],
+                ele["resized_width"],
+                factor=image_factor,
+            )
+        else:
+            resized_height, resized_width = smart_resize(
+                height,
+                width,
+                factor=image_factor,
+                min_pixels=min_pixels,
+                max_pixels=max_pixels,
+            )
+        video = transforms.functional.resize(
+            video,
+            [resized_height, resized_width],
+            interpolation=InterpolationMode.BICUBIC,
+            antialias=True,
+        ).float()
+        return video, total_dur, nframes, second_per_grid
+    else:
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        images = [
+            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)
+        if len(images) < nframes:
+            images.extend([images[-1]] * (nframes - len(images)))
+        return images, None, None, None
+
+
+def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations = [conversations]
+    for conversation in conversations:
+        for message in conversation:
+            if isinstance(message["content"], list):
+                for ele in message["content"]:
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele["type"] in ("image", "image_url", "video")
+                    ):
+                        vision_infos.append(ele)
+    return vision_infos
+
+
+def process_vision_info(
+    conversations: list[dict] | list[list[dict]],
+) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]:
+    vision_infos = extract_vision_info(conversations)
+    # Read images or videos
+    image_inputs = []
+    video_inputs = []
+    for vision_info in vision_infos:
+        if "image" in vision_info or "image_url" in vision_info:
+            image_inputs.append(fetch_image(vision_info))
+        elif "video" in vision_info:
+            video_inputs.append(fetch_video(vision_info))
+        else:
+            raise ValueError("image, image_url or video should in content.")
+    if len(image_inputs) == 0:
+        image_inputs = None
+    if len(video_inputs) == 0:
+        video_inputs = None
+    return image_inputs, video_inputs

examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh

@@ -0,0 +1,8 @@
+python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --txt-prompts top100.txt

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

@@ -0,0 +1,8 @@
+python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."

examples/offline_inference/qwen2_5_omni/utils.py

@@ -0,0 +1,312 @@
+import tempfile
+from typing import Optional, Union
+from urllib.request import urlopen
+
+import librosa
+import requests
+import resampy
+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"] = {}
+
+
+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 resample_wav_to_16khz(input_filepath):
+    data, original_sample_rate = sf.read(input_filepath)
+    # Only use the first channel
+    if len(data.shape) > 1:
+        data = data[:, 0]
+    # resample to 16kHz
+    data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000)
+    return data_resampled
+
+
+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,
+            start_pts=0.0,
+            end_pts=None,
+            pts_unit="sec",
+            output_format="TCHW",
+        )
+
+        total_frames, video_fps = video.size(0), info["video_fps"]
+        total_duration = round(total_frames / video_fps, 3)
+        nframes = int(total_frames / video_fps * fps)
+
+        frame_timestamps = total_duration * torch.arange(1, nframes + 1) / nframes
+        grid_timestamps = frame_timestamps[::2]
+        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
+
+        idx = torch.linspace(0, video.size(0) - 1, nframes).round().long()
+        video = video[idx]
+
+        if args.legacy_omni_video:
+            return [video, total_duration, nframes, second_per_grid.item()]
+        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})
+        if total_duration is None and nframes is None:
+            nframes = len(video)
+            total_duration = 0.5 * nframes
+            second_per_grid = 1.0
+        if args.legacy_omni_video:
+            return [video, total_duration, nframes, second_per_grid]
+        else:
+            return video
+
+    def read_video(video_file_name: str):
+        if args.use_torchvision:
+            return read_video_with_torchvision(video_file_name)
+        else:
+            return read_video_with_transformers(video_file_name)
+
+    if isinstance(video_url, str) and video_url.startswith("http"):
+        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}, status_code:{resp.status_code}, resp:{resp}"
+            )
+
+            temp_video_file.write(urlopen(video_url).read())
+            temp_video_file_path = temp_video_file.name
+            video_file_name = temp_video_file_path
+            return read_video(video_file_name)
+    else:
+        video_file_name = video_url
+        return read_video(video_file_name)
+
+
+def make_inputs_qwen2_omni(
+    args,
+    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
+    if args.model not in _PROCESSOR_CACHE:
+        _PROCESSOR_CACHE[args.model] = AutoProcessor.from_pretrained(args.model)
+    processor = _PROCESSOR_CACHE[args.model]
+
+    config = _CONFIG_CACHE.get(args.model)
+    if config is None:
+        try:
+            config = AutoConfig.from_pretrained(args.model)
+        except Exception:
+            config = None
+        _CONFIG_CACHE[args.model] = config  # cache even if None to avoid retry storms
+
+    # 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 = "Qwen2_5OmniModel" not in config.architectures
+        else:
+            args.legacy_omni_video = True
+
+    audios, images, videos = [], [], []
+    for message in messages:
+        if not isinstance(message["content"], list):
+            message["content"] = [
+                {
+                    "type": "text",
+                    "text": message["content"],
+                }
+            ]
+        index, num_contents = 0, len(message["content"])
+        while index < num_contents:
+            ele = message["content"][index]
+            if "type" not in ele:
+                if "text" in ele:
+                    ele["type"] = "text"
+                elif "audio" in ele:
+                    ele["type"] = "audio"
+                elif "audio_url" in ele:
+                    ele["type"] = "audio_url"
+                elif "image" in ele:
+                    ele["type"] = "image"
+                elif "image_url" in ele:
+                    ele["type"] = "image_url"
+                elif "video" in ele:
+                    ele["type"] = "video"
+                elif "video_url" in ele:
+                    ele["type"] = "video_url"
+                else:
+                    raise ValueError(f"Unknown ele: {ele}")
+
+            if ele["type"] == "audio" or ele["type"] == "audio_url":
+                if "audio_url" in ele:
+                    audio_key = "audio_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_audio_file:
+                        temp_audio_file.write(urlopen(ele[audio_key]).read())
+                        temp_audio_file_path = temp_audio_file.name
+                        audios.append(resample_wav_to_16khz(temp_audio_file_path))
+                        ele["audio"] = temp_audio_file_path
+                elif "audio" in ele:
+                    audio_key = "audio"
+                    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"):
+                # use video as audio as well
+                if "video_url" in ele:
+                    audio_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele[audio_key]).read())
+                        temp_video_file_path = temp_video_file.name
+                        ele[audio_key] = temp_video_file_path
+                        audios.append(librosa.load(temp_video_file_path, sr=16000)[0])
+                        videos.append(fetch_and_read_video(args, temp_video_file_path))
+                        ele["video"] = temp_video_file_path
+                elif "video" in ele:
+                    audio_key = "video"
+                    audios.append(librosa.load(ele[audio_key], sr=16000)[0])
+                    videos.append(fetch_and_read_video(args, audio_key))
+                else:
+                    raise ValueError(f"Unknown ele {ele}")
+                # insert a audio after the video
+                message["content"].insert(
+                    index + 1,
+                    {
+                        "type": "audio",
+                        "audio": ele[audio_key],
+                    },
+                )
+                # no need to load the added audio again
+                index += 1
+            elif ele["type"] == "video" or ele["type"] == "video_url":
+                if "video_url" in ele:
+                    video_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele["video_url"]).read())
+                        temp_video_file_path = temp_video_file.name
+                        videos.append(fetch_and_read_video(args, temp_video_file))
+                        ele["video"] = temp_video_file_path
+                else:
+                    video_key = "video"
+                    videos.append(fetch_and_read_video(args, ele[video_key]))
+            elif ele["type"] == "image" or ele["type"] == "image_url":
+                images.append(fetch_image(ele))
+
+            # move to the next content
+            index += 1
+
+    prompt = processor.apply_chat_template(
+        messages,
+        tokenize=tokenize,
+        add_generation_prompt=True,
+        add_vision_id=True,
+    )
+
+    audios = audios if len(audios) > 0 else None
+    images = images if len(images) > 0 else None
+    videos = videos if len(videos) > 0 else None
+
+    multi_modal_data = {}
+    if audios:
+        multi_modal_data["audio"] = audios
+    if images:
+        multi_modal_data["image"] = images
+    if videos:
+        multi_modal_data["video"] = videos
+
+    if isinstance(prompt, list) and isinstance(prompt[0], (list, str)):
+        prompt = prompt[0]
+
+    if tokenize:
+        return OmniTokensPrompt(
+            prompt_token_ids=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+    else:
+        return TextPrompt(
+            prompt=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+
+
+def make_text_prompt(args, prompt):
+    messages = [
+        get_system_prompt(),
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt},
+            ],
+        },
+    ]
+
+    prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize)
+    return prompt
+
+
+def make_audio_in_video_v2_prompt(args):
+    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": "video_url",
+                    "video_url": ("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"),
+                },
+            ],
+        },
+    ]
+    prompt = make_inputs_qwen2_omni(
+        args,
+        messages,
+        use_audio_in_video=True,
+        tokenize=args.tokenize,
+    )
+    return prompt
+
+
+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":
+        prompt = make_audio_in_video_v2_prompt(args)
+    else:
+        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
+    return prompt

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

@@ -0,0 +1,157 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "vllm-omni"
+version = "0.11.0rc1"
+description = "vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures"
+readme = "README.md"
+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 = "vllm-project@gmail.com"}
+]
+keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "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 = [
+    # 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",
+    "mypy>=1.0.0",
+    "pre-commit>=3.0.0",
+]
+
+docs = [
+    "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/vllm-project/vllm-omni"
+Repository = "https://github.com/vllm-project/vllm-omni"
+Documentation = "https://vllm-omni.readthedocs.io"
+"Bug Tracker" = "https://github.com/vllm-project/vllm-omni/issues"
+
+[project.scripts]
+vllm = "vllm_omni.entrypoints.cli.main:main"
+vllm-omni = "vllm_omni.entrypoints.cli.main:main"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["vllm_omni*"]
+
+[tool.ruff]
+line-length = 120
+exclude = [
+    ".eggs",
+    ".git",
+    ".hg",
+    ".mypy_cache",
+    ".tox",
+    ".venv",
+    "build",
+    "dist",
+    "vllm_omni.egg-info",
+]
+
+[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.ruff.lint.per-file-ignores]
+"examples/**" = ["E501"]  # Allow long lines in examples
+"tests/**" = ["E501"]  # Allow long lines in tests
+
+[tool.mypy]
+python_version = "3.12, 3.13"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+check_untyped_defs = true
+disallow_untyped_decorators = true
+no_implicit_optional = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+warn_no_return = true
+warn_unreachable = true
+strict_equality = true
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py", "*_test.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+addopts = [
+    "--strict-markers",
+    "--strict-config",
+    "--cov=vllm_omni",
+    "--cov-report=term-missing",
+    "--cov-report=html",
+    "--cov-report=xml",
+]
+markers = [
+    "unit: Unit tests",
+    "integration: Integration tests",
+    "benchmark: Benchmark tests",
+    "slow: Slow tests",
+]
+
+[tool.typos.default]
+extend-ignore-identifiers-re = [
+    ".*_thw",
+    ".*thw",
+    "ein",
+    ".*arange"
+]

scripts/README.md

@@ -0,0 +1,108 @@
+# vLLM-omni Testing Scripts
+
+This directory contains automated testing scripts for vLLM-omni serving functionality.
+
+## Scripts Overview
+
+### 1. `test_serving.sh` - Comprehensive Test Suite
+Full-featured testing script that performs complete validation of the vLLM-omni serving functionality.
+
+**Features:**
+- Model existence validation
+- Environment setup verification
+- Import functionality testing
+- Server startup and health checks
+- Text generation testing
+- Performance benchmarking
+- API client integration testing
+- Comprehensive logging and error handling
+
+**Usage:**
+```bash
+# Use default model and port
+./scripts/test_serving.sh
+
+# Use specific model and port
+./scripts/test_serving.sh ./models/Qwen3-0.6B 8000
+
+# Use HuggingFace model
+./scripts/test_serving.sh Qwen/Qwen3-0.6B 8001
+
+# Show help
+./scripts/test_serving.sh --help
+```
+
+**Test Coverage:**
+- ✅ Model loading and server startup
+- ✅ Health and info endpoints
+- ✅ Text generation functionality
+- ✅ Performance metrics
+- ✅ API client integration
+- ✅ All imports working correctly
+
+### 2. `quick_test.sh` - Fast Validation
+Lightweight script for quick validation after making changes.
+
+**Features:**
+- Fast import testing
+- Basic server startup
+- Health endpoint validation
+- Simple text generation test
+- Retry mechanism for reliability
+
+**Usage:**
+```bash
+# Use default port (8000)
+./scripts/quick_test.sh
+
+# Use specific port
+./scripts/quick_test.sh 8001
+```
+
+**Test Coverage:**
+- ✅ Import functionality
+- ✅ Server startup
+- ✅ Health endpoint
+- ✅ Basic text generation
+
+## Prerequisites
+
+1. **Conda Environment**: Ensure `vllm_omni` environment is activated
+2. **Model Available**: Qwen3-0.6B model should be available in `./models/Qwen3-0.6B/`
+3. **Dependencies**: All vLLM-omni dependencies should be installed
+
+## Environment Setup
+
+```bash
+# Activate conda environment
+conda activate vllm_omni
+
+# Verify installation
+python -c "import vllm_omni; print('vLLM-omni ready')"
+```
+
+## Model Setup
+
+The scripts expect the Qwen3-0.6B model to be available. You can:
+
+1. **Use local model**: Place model in `./models/Qwen3-0.6B/`
+2. **Use HuggingFace model**: Pass `Qwen/Qwen3-0.6B` as model path
+3. **Download model**: Use the `download_models.py` script
+
+```bash
+# Download model using the provided script
+python scripts/download_models.py
+```
+
+## Output and Logging
+
+### Test Results
+- **Success**: Green `[SUCCESS]` messages
+- **Info**: Blue `[INFO]` messages  
+- **Warnings**: Yellow `[WARNING]` messages
+- **Errors**: Red `[ERROR]` messages
+
+### Log Files
+- `server.log`: Server startup and runtime logs
+- `api_client_test.log`: API client test results
+- `simple_usage_test.log`: Simple usage test results

scripts/test_serving.sh

@@ -0,0 +1,363 @@
+#!/bin/bash
+
+# vLLM-omni Serving Functionality Test Script
+# This script tests the complete serving functionality of vLLM-omni
+# Usage: ./scripts/test_serving.sh [model_path] [port]
+
+set -e  # Exit on any error
+
+# Configuration
+MODEL_PATH=${1:-"./models/Qwen3-0.6B"}
+PORT=${2:-8000}
+HOST="localhost"
+TIMEOUT=30
+SERVER_STARTUP_WAIT=15
+
+# Colors for output
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m' # No Color
+
+# Logging functions
+log_info() {
+    echo -e "${BLUE}[INFO]${NC} $1"
+}
+
+log_success() {
+    echo -e "${GREEN}[SUCCESS]${NC} $1"
+}
+
+log_warning() {
+    echo -e "${YELLOW}[WARNING]${NC} $1"
+}
+
+log_error() {
+    echo -e "${RED}[ERROR]${NC} $1"
+}
+
+# Cleanup function
+cleanup() {
+    log_info "Cleaning up..."
+    if [ ! -z "$SERVER_PID" ]; then
+        log_info "Stopping server (PID: $SERVER_PID)..."
+        kill $SERVER_PID 2>/dev/null || true
+        wait $SERVER_PID 2>/dev/null || true
+    fi
+    # Kill any remaining vllm processes
+    pkill -f "vllm serve" 2>/dev/null || true
+    log_info "Cleanup completed"
+}
+
+# Set up trap for cleanup on exit
+trap cleanup EXIT
+
+# Check if model exists
+check_model() {
+    log_info "Checking if model exists at: $MODEL_PATH"
+    if [ ! -d "$MODEL_PATH" ]; then
+        log_error "Model directory not found: $MODEL_PATH"
+        log_info "Available models:"
+        ls -la models/ 2>/dev/null || log_warning "No models directory found"
+        exit 1
+    fi
+    log_success "Model found: $MODEL_PATH"
+}
+
+# Check if conda environment is activated
+check_environment() {
+    log_info "Checking conda environment..."
+    if [[ "$CONDA_DEFAULT_ENV" != "vllm_omni" ]]; then
+        log_warning "vllm_omni conda environment not activated"
+        log_info "Attempting to activate vllm_omni environment..."
+        source $(conda info --base)/etc/profile.d/conda.sh
+        conda activate vllm_omni
+        if [[ "$CONDA_DEFAULT_ENV" != "vllm_omni" ]]; then
+            log_error "Failed to activate vllm_omni environment"
+            exit 1
+        fi
+    fi
+    log_success "Environment check passed: $CONDA_DEFAULT_ENV"
+}
+
+# Test import functionality
+test_imports() {
+    log_info "Testing imports..."
+    python -c "
+import vllm_omni
+from vllm_omni.entrypoints.omni_llm import OmniLLM, AsyncOmniLLM
+print('✅ All imports successful')
+" || {
+        log_error "Import test failed"
+        exit 1
+    }
+    log_success "Import test passed"
+}
+
+# Start the server
+start_server() {
+    log_info "Starting vLLM-omni server on port $PORT..."
+    log_info "Command: vllm serve $MODEL_PATH --omni --port $PORT"
+
+    # Start server in background
+    vllm serve "$MODEL_PATH" --omni --port "$PORT" > server.log 2>&1 &
+    SERVER_PID=$!
+
+    log_info "Server started with PID: $SERVER_PID"
+    log_info "Waiting $SERVER_STARTUP_WAIT seconds for server to initialize..."
+    sleep $SERVER_STARTUP_WAIT
+
+    # Check if server is still running
+    if ! kill -0 $SERVER_PID 2>/dev/null; then
+        log_error "Server failed to start"
+        log_info "Server logs:"
+        cat server.log
+        exit 1
+    fi
+
+    log_success "Server appears to be running"
+}
+
+# Test health endpoint
+test_health() {
+    log_info "Testing health endpoint..."
+    for attempt in {1..10}; do
+        local response=$(curl -s -w "%{http_code}" http://$HOST:$PORT/health 2>/dev/null || echo "000")
+        local http_code="${response: -3}"
+        local body="${response%???}"
+
+        if [ "$http_code" = "200" ]; then
+            log_success "Health check passed: $body"
+            return 0
+        fi
+
+        log_warning "Health check attempt $attempt failed (HTTP $http_code). Retrying in 3s..."
+        sleep 3
+    done
+
+    log_error "Health check failed after multiple attempts"
+    return 1
+}
+
+# Test info endpoint
+test_info() {
+    log_info "Testing info endpoint..."
+    local response=$(curl -s -w "%{http_code}" http://$HOST:$PORT/info 2>/dev/null || echo "000")
+    local http_code="${response: -3}"
+    local body="${response%???}"
+
+    if [ "$http_code" = "200" ]; then
+        log_success "Info endpoint working"
+        echo "$body" | python -m json.tool 2>/dev/null || log_warning "Info response not valid JSON"
+    else
+        log_error "Info endpoint failed (HTTP $http_code): $body"
+        return 1
+    fi
+}
+
+# Test text generation
+test_generation() {
+    log_info "Testing text generation..."
+
+    # Test 1: Simple generation
+    local response=$(curl -s -X POST http://$HOST:$PORT/generate \
+        -H "Content-Type: application/json" \
+        -d '{"prompts": ["Test the server functionality"], "max_tokens": 20, "temperature": 0.7}' \
+        2>/dev/null || echo "{}")
+
+    if echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); exit(0 if 'outputs' in data and len(data['outputs']) > 0 else 1)" 2>/dev/null; then
+        log_success "Text generation test passed"
+        echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); print('Generated text:', data['outputs'][0]['text'][:100] + '...' if len(data['outputs'][0]['text']) > 100 else data['outputs'][0]['text'])" 2>/dev/null
+    else
+        log_error "Text generation test failed"
+        echo "Response: $response"
+        return 1
+    fi
+}
+
+# Test API client example
+test_api_client() {
+    log_info "Testing API client example..."
+    if [ -f "examples/basic/api_client.py" ]; then
+        python examples/basic/api_client.py > api_client_test.log 2>&1 || {
+            log_error "API client test failed"
+            log_info "API client logs:"
+            cat api_client_test.log
+            return 1
+        }
+        log_success "API client test passed"
+    else
+        log_warning "API client example not found, skipping"
+    fi
+}
+
+# Test AR → DiT pipeline example
+test_ar_dit_pipeline() {
+    log_info "Testing AR → DiT diffusers pipeline example..."
+
+    if ! python - 2>/dev/null <<'PY'
+import importlib
+for pkg in ("torch", "diffusers"):
+    if importlib.util.find_spec(pkg) is None:
+        raise ImportError(pkg)
+PY
+    then
+        log_warning "Required packages for diffusers pipeline not found, skipping AR → DiT test"
+        return 0
+    fi
+
+    if [ ! -f "examples/omni/ar_dit_diffusers.py" ]; then
+        log_warning "AR → DiT example not found, skipping"
+        return 0
+    fi
+
+    local output_path="logs/ar_dit_pipeline.png"
+    local log_path="logs/ar_dit_pipeline.log"
+    rm -f "$output_path"
+
+    local cmd=(python examples/omni/ar_dit_diffusers.py
+        --prompt "Smoke test prompt of a lighthouse"
+        --seed 0
+        --output "$output_path")
+
+    local env_prefix=()
+    if [ -n "$AR_MODEL" ]; then
+        env_prefix+=("AR_MODEL=$AR_MODEL")
+    fi
+    if [ -n "$DIT_MODEL" ]; then
+        env_prefix+=("DIT_MODEL=$DIT_MODEL")
+    fi
+
+    if ! env "${env_prefix[@]}" "${cmd[@]}" >"$log_path" 2>&1; then
+        log_error "AR → DiT pipeline test failed"
+        log_info "AR → DiT logs:"
+        cat "$log_path"
+        return 1
+    fi
+
+    if [ ! -f "$output_path" ]; then
+        log_error "AR → DiT pipeline did not produce an image"
+        log_info "AR → DiT logs:"
+        cat "$log_path"
+        return 1
+    fi
+
+    log_success "AR → DiT pipeline test passed (output: $output_path)"
+}
+
+# Test simple usage example
+test_simple_usage() {
+    log_info "Testing simple usage example..."
+    if [ -f "examples/basic/simple_usage.py" ]; then
+        # Run with timeout to avoid hanging
+        timeout 60 python examples/basic/simple_usage.py > simple_usage_test.log 2>&1 || {
+            local exit_code=$?
+            if [ $exit_code -eq 124 ]; then
+                log_warning "Simple usage test timed out (60s), but this might be normal for model loading"
+            else
+                log_error "Simple usage test failed (exit code: $exit_code)"
+                log_info "Simple usage logs:"
+                cat simple_usage_test.log
+                return 1
+            fi
+        }
+        log_success "Simple usage test passed"
+    else
+        log_warning "Simple usage example not found, skipping"
+    fi
+}
+
+# Performance test
+test_performance() {
+    log_info "Running performance test..."
+
+    local start_time=$(date +%s)
+    local response=$(curl -s -X POST http://$HOST:$PORT/generate \
+        -H "Content-Type: application/json" \
+        -d '{"prompts": ["Performance test prompt"], "max_tokens": 50, "temperature": 0.7}' \
+        2>/dev/null)
+    local end_time=$(date +%s)
+    local duration=$((end_time - start_time))
+
+    if echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); exit(0 if 'outputs' in data else 1)" 2>/dev/null; then
+        log_success "Performance test passed (${duration}s response time)"
+    else
+        log_error "Performance test failed"
+        return 1
+    fi
+}
+
+# Main test function
+run_tests() {
+    log_info "Starting vLLM-omni serving functionality tests..."
+    log_info "Model: $MODEL_PATH"
+    log_info "Port: $PORT"
+    log_info "Host: $HOST"
+    echo "=========================================="
+
+    # Run all tests
+    check_model
+    check_environment
+    test_imports
+    start_server
+
+    # Wait a bit more for server to be fully ready
+    sleep 5
+
+    test_health
+    test_info
+    test_generation
+    test_performance
+    test_api_client
+    test_ar_dit_pipeline
+
+    # Note: Simple usage test is commented out as it starts its own server
+    # test_simple_usage
+
+    log_success "All tests completed successfully!"
+    echo "=========================================="
+    log_info "Test Summary:"
+    log_info "✅ Model loading and server startup"
+    log_info "✅ Health and info endpoints"
+    log_info "✅ Text generation functionality"
+    log_info "✅ Performance metrics"
+    log_info "✅ API client integration"
+    log_info "✅ AR → DiT diffusers pipeline"
+    log_info "✅ All imports working correctly"
+}
+
+# Show usage
+show_usage() {
+    echo "Usage: $0 [model_path] [port]"
+    echo ""
+    echo "Arguments:"
+    echo "  model_path    Path to the model directory (default: ./models/Qwen3-0.6B)"
+    echo "  port          Port to run the server on (default: 8000)"
+    echo ""
+    echo "Examples:"
+    echo "  $0                                    # Use default model and port"
+    echo "  $0 ./models/Qwen3-0.6B 8001          # Use specific model and port"
+    echo "  $0 Qwen/Qwen3-0.6B 8000              # Use HuggingFace model"
+}
+
+# Main execution
+main() {
+    # Check for help flag
+    if [[ "$1" == "-h" || "$1" == "--help" ]]; then
+        show_usage
+        exit 0
+    fi
+
+    # Change to script directory
+    cd "$(dirname "$0")/.."
+
+    # Create logs directory if it doesn't exist
+    mkdir -p logs
+
+    # Run tests
+    run_tests
+}
+
+# Run main function with all arguments
+main "$@"

tests/__init__.py

@@ -0,0 +1,6 @@
+"""
+Test suite for vLLM-omni.
+
+This package contains unit tests, integration tests, and benchmarks
+for vLLM-omni.
+"""

vllm_omni/__init__.py

@@ -0,0 +1,41 @@
+"""
+vLLM-omni: Multi-modality models inference and serving with
+non-autoregressive structures.
+
+This package extends vLLM beyond traditional text-based, autoregressive
+generation to support multi-modality models with non-autoregressive
+structures and non-textual outputs.
+
+Architecture:
+- 🟡 Modified: vLLM components modified for multimodal support
+- 🔴 Added: New components for multimodal and non-autoregressive
+  processing
+"""
+
+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",
+    # Configuration
+    "OmniModelConfig",
+    # All other components are available through their respective modules
+    # processors.*, schedulers.*, executors.*, etc.
+]

vllm_omni/config/__init__.py

@@ -0,0 +1,9 @@
+"""
+Configuration module for vLLM-omni.
+"""
+
+from vllm_omni.config.model import OmniModelConfig
+
+__all__ = [
+    "OmniModelConfig",
+]

vllm_omni/config/model.py

@@ -0,0 +1,27 @@
+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
+
+
+@config
+@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
+class OmniModelConfig(ModelConfig):
+    """Configuration for Omni models, extending the base 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
+    def registry(self):
+        return me_models.OmniModelRegistry
+
+    @property
+    def architectures(self) -> list[str]:
+        return [self.model_arch]

vllm_omni/core/dit_cache_manager.py

@@ -0,0 +1,13 @@
+"""
+DiT Cache Manager for vLLM-omni.
+
+This module provides caching functionality for DiT (Diffusion Transformer) models
+to optimize inference performance and memory usage.
+"""
+
+
+class DiTCacheManager:
+    """Manages DiT-specific caching for optimized inference."""
+
+    def __init__(self, config):
+        pass

vllm_omni/core/sched/__init__.py

@@ -0,0 +1,13 @@
+"""
+Scheduling components for vLLM-omni.
+"""
+
+from .diffusion_scheduler import DiffusionScheduler
+from .output import OmniNewRequestData
+from .scheduler import OmniScheduler
+
+__all__ = [
+    "OmniScheduler",
+    "DiffusionScheduler",
+    "OmniNewRequestData",
+]

vllm_omni/core/sched/diffusion_scheduler.py

@@ -0,0 +1,299 @@
+import time
+from collections import defaultdict
+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.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
+
+
+class DiffusionScheduler(OmniScheduler):
+    def schedule(self) -> SchedulerOutput:
+        """Diffusion fast path:
+        - Feed all input tokens of the request at once
+          (if 0, allocate 1 placeholder token).
+        - If the token budget cannot be satisfied at once, fall back to the
+          default vLLM scheduling.
+        """
+
+        # Select requests with zero prompt and using pooling
+        # (diffusion results returned via pooler_output)
+        token_budget = self.max_num_scheduled_tokens
+        capacity = self.max_num_running_reqs - len(self.running)
+        scheduled_timestamp = time.monotonic()
+
+        scheduled_new_reqs: list[Request] = []
+        scheduled_resumed_reqs: list[Request] = []
+        scheduled_running_reqs: list[Request] = []
+
+        req_to_new_blocks: dict[str, KVCacheBlocks] = {}
+        num_scheduled_tokens: dict[str, int] = {}
+        scheduled_spec_decode_tokens: dict[str, list[int]] = {}
+        scheduled_encoder_inputs: dict[str, list[int]] = {}
+
+        # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
+        skipped_waiting_requests = create_request_queue(self.policy)
+
+        # Fast path selection and scheduling (treat all as diffusion requests,
+        # independent of pooling_params)
+        while self.waiting and token_budget > 0 and capacity > 0:
+            request = self.waiting.peek_request()
+            # Uniformly treat as diffusion. A feature flag can be added later
+            # via config or request tag.
+
+            # Allocate all input tokens for the request in one shot
+            # (allocate 1 placeholder if zero)
+            required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
+            if required_tokens > token_budget:
+                # Insufficient budget to process all inputs at once;
+                # stop fast path attempt
+                break
+            num_new_tokens = required_tokens
+            new_blocks = self.kv_cache_manager.allocate_slots(
+                request,
+                num_new_tokens,
+                num_lookahead_tokens=self.num_lookahead_tokens,
+            )
+            if new_blocks is None:
+                # Allocation failed (e.g., VRAM pressure); stop fast path and
+                # fall back to default scheduling
+                # Put the current request back to the head of the waiting queue
+                # Note: the original queue order is preserved
+                break
+
+            # Officially schedule this request
+            request = self.waiting.pop_request()
+            self.running.append(request)
+            request.status = RequestStatus.RUNNING
+            if self.log_stats:
+                request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
+
+            req_to_new_blocks[request.request_id] = new_blocks
+            num_scheduled_tokens[request.request_id] = num_new_tokens
+            token_budget -= num_new_tokens
+            capacity -= 1
+            scheduled_new_reqs.append(request)
+
+        # Return skipped waiting requests
+        if skipped_waiting_requests:
+            self.waiting.prepend_requests(skipped_waiting_requests)
+
+        # If fast path scheduled none, fall back to the original scheduling
+        if not num_scheduled_tokens:
+            return super().schedule()
+
+        # Compute common prefix blocks (aligned with v1)
+        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)
+            )
+
+        # Assemble SchedulerOutput
+        new_reqs_data = [
+            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(
+            scheduled_running_reqs,
+            scheduled_resumed_reqs,
+            num_scheduled_tokens,
+            scheduled_spec_decode_tokens,
+            req_to_new_blocks,
+        )
+        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
+        )
+
+        total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=new_reqs_data,
+            scheduled_cached_reqs=cached_reqs_data,
+            num_scheduled_tokens=num_scheduled_tokens,
+            total_num_scheduled_tokens=total_num_scheduled_tokens,
+            scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
+            scheduled_encoder_inputs=scheduled_encoder_inputs,
+            num_common_prefix_blocks=num_common_prefix_blocks,
+            finished_req_ids=self.finished_req_ids,
+            free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
+            structured_output_request_ids=structured_output_request_ids,
+            grammar_bitmask=grammar_bitmask,
+        )
+
+        # KVTransfer: package metadata
+        if self.connector is not None:
+            meta = self.connector.build_connector_meta(scheduler_output)
+            scheduler_output.kv_connector_metadata = meta
+
+        # Update internal state (advance num_computed_tokens, free encoder inputs,
+        # etc.)
+        self._update_after_schedule(scheduler_output)
+        return scheduler_output
+
+    """
+    Scheduler for the diffusion model.
+    This scheduler is modified to stop the request immediately for the diffusion model.
+    This is because the diffusion model can generate the final image/audio in one step.
+    Note: This is just a minimal modification to the original scheduler,
+    and there should be some further efforts to optimize the scheduler.
+    The original scheduler is still used for the AR model.
+    """
+
+    def update_from_output(
+        self,
+        scheduler_output: SchedulerOutput,
+        model_runner_output: OmniModelRunnerOutput,
+    ) -> dict[int, EngineCoreOutputs]:
+        """Update the scheduler state based on the model runner output.
+
+        This method is modified to stop the request immediately for the diffusion model.
+        """
+        sampled_token_ids = model_runner_output.sampled_token_ids
+        logprobs = model_runner_output.logprobs
+        prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
+        num_scheduled_tokens = scheduler_output.num_scheduled_tokens
+        pooler_outputs = model_runner_output.pooler_output
+        num_nans_in_logits = model_runner_output.num_nans_in_logits
+        kv_connector_output = model_runner_output.kv_connector_output
+
+        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
+
+        # 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
+        # to avoid expensive operations inside the loop.
+        stopped_running_reqs: set[Request] = set()
+        stopped_preempted_reqs: set[Request] = set()
+        for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
+            assert num_tokens_scheduled > 0
+            request = self.requests.get(req_id)
+            if request is None:
+                # The request is already finished. This can happen if the
+                # request is aborted while the model is executing it (e.g.,
+                # in pipeline parallelism).
+                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 []
+
+            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
+                num_rejected = num_draft_tokens - num_accepted
+                # num_computed_tokens represents the number of tokens
+                # processed in the current step, considering scheduled
+                # tokens and rejections. If some tokens are rejected,
+                # num_computed_tokens is decreased by the number of rejected
+                # tokens.
+                request.num_computed_tokens -= num_rejected
+                spec_decoding_stats = self.make_spec_decoding_stats(
+                    spec_decoding_stats,
+                    num_draft_tokens=num_draft_tokens,
+                    num_accepted_tokens=num_accepted,
+                )
+
+            new_logprobs = None
+            new_token_ids = generated_token_ids
+            kv_transfer_params = None
+            status_before_stop = request.status
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]
+
+            # Diffusion request: completes in one step; mark finished and free resources
+            request.status = RequestStatus.FINISHED_STOPPED
+            # Optional: set a stop_reason for front-end clarity
+            # (does not affect protocol)
+            request.stop_reason = request.stop_reason or "diffusion_done"
+            kv_transfer_params = self._free_request(request)
+            if status_before_stop == RequestStatus.RUNNING:
+                stopped_running_reqs.add(request)
+            else:
+                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:
+                # 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)
+
+            if new_token_ids and self.structured_output_manager.should_advance(request):
+                # 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(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:
+                request.num_nans_in_logits = num_nans_in_logits[req_id]
+
+            # 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(
+                        request_id=req_id,
+                        new_token_ids=new_token_ids,
+                        finish_reason=request.get_finished_reason(),
+                        new_logprobs=new_logprobs,
+                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                        pooling_output=pooler_output,
+                        stop_reason=request.stop_reason,
+                        events=request.take_events(),
+                        kv_transfer_params=kv_transfer_params,
+                        trace_headers=request.trace_headers,
+                        num_cached_tokens=request.num_cached_tokens,
+                    )
+                )
+            else:
+                # Invariant: EngineCore returns no partial prefill outputs.
+                assert not prompt_logprobs_tensors
+
+        # Remove the stopped requests from the running and waiting queues.
+        if stopped_running_reqs:
+            if stopped_running_reqs:
+                self.running = remove_all(self.running, stopped_running_reqs)
+        if stopped_preempted_reqs:
+            # This is a rare case and unlikely to impact performance.
+            self.waiting.remove_requests(stopped_preempted_reqs)
+
+        # KV Connector: update state for finished KV Transfers.
+        if model_runner_output.kv_connector_output:
+            self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output)
+
+        # 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()}
+
+        finished_req_ids = self.finished_req_ids_dict
+        if finished_req_ids:
+            # Include ids of requests that finished since last outputs
+            # were sent.
+            for client_index, finished_set in finished_req_ids.items():
+                # Set finished request set in EngineCoreOutputs for this client.
+                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)
+            finished_req_ids.clear()
+
+        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
+                # outputs this step.
+                engine_core_outputs[0] = eco = EngineCoreOutputs()
+            eco.scheduler_stats = stats
+
+        return engine_core_outputs

vllm_omni/core/sched/output.py

@@ -0,0 +1,34 @@
+from dataclasses import dataclass
+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
+
+
+@dataclass
+class OmniNewRequestData(NewRequestData):
+    # Optional serialized prompt embeddings
+    prompt_embeds: Optional[PromptEmbedsPayload] = None
+    # Optional serialized additional information
+    additional_information: Optional[AdditionalInformationPayload] = None
+
+    @classmethod
+    def from_request(
+        cls,
+        request: Request,
+        block_ids: tuple[list[int], ...],
+    ) -> "OmniNewRequestData":
+        return cls(
+            req_id=request.request_id,
+            prompt_token_ids=request.prompt_token_ids,
+            mm_features=request.mm_features,
+            sampling_params=request.sampling_params,
+            pooling_params=request.pooling_params,
+            block_ids=block_ids,
+            num_computed_tokens=request.num_computed_tokens,
+            lora_request=request.lora_request,
+            prompt_embeds=request.prompt_embeds,
+            additional_information=request.additional_information,
+        )

vllm_omni/core/sched/scheduler.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
+
+
+class OmniScheduler(VLLMScheduler):
+    """
+    OmniScheduler: Scheduler for vLLM-omni multimodal processing.
+
+    This scheduler extends vLLM's scheduler to support multimodal and
+    non-autoregressive processing with additional fields and methods
+    specific to vLLM-omni.
+    """
+
+    # Ensure scheduled_new_reqs carry omni-specific payloads
+    # (e.g., additional_information)
+    def schedule(self) -> SchedulerOutput:  # type: ignore[override]
+        scheduler_output = super().schedule()
+        try:
+            # Late import to avoid circulars in some launch modes
+            from .output import OmniNewRequestData
+
+            # Rewrap base NewRequestData entries with OmniNewRequestData,
+            # enriching with request-level payloads
+            new_list = []
+            for nr in scheduler_output.scheduled_new_reqs:
+                req_id = getattr(nr, "req_id", None)
+                request = self.requests.get(req_id) if req_id else None
+                # Build omni entry preserving all base fields
+                omni_nr = OmniNewRequestData(
+                    req_id=nr.req_id,
+                    prompt_token_ids=nr.prompt_token_ids,
+                    mm_features=nr.mm_features,
+                    sampling_params=nr.sampling_params,
+                    pooling_params=nr.pooling_params,
+                    block_ids=nr.block_ids,
+                    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),
+                )
+                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")
+
+        return scheduler_output

vllm_omni/engine/__init__.py

@@ -0,0 +1,55 @@
+"""
+Engine components for vLLM-omni.
+"""
+
+from typing import Any, Optional
+
+import msgspec
+from vllm.v1.engine import EngineCoreRequest
+
+
+class PromptEmbedsPayload(msgspec.Struct):
+    """Serialized prompt embeddings payload for direct transfer.
+
+    data: raw bytes of the tensor in row-major order
+    shape: [seq_len, hidden_size]
+    dtype: torch dtype name (e.g., "float16", "float32")
+    """
+
+    data: bytes
+    shape: list[int]
+    dtype: str
+
+
+class AdditionalInformationEntry(msgspec.Struct):
+    """One entry of additional_information.
+
+    Two supported forms are encoded:
+      - tensor: data/shape/dtype
+      - list: a Python list (msgspec-serializable)
+    Exactly one of (tensor_data, list_data) should be non-None.
+    """
+
+    # Tensor form
+    tensor_data: Optional[bytes] = None
+    tensor_shape: Optional[list[int]] = None
+    tensor_dtype: Optional[str] = None
+
+    # List form
+    list_data: Optional[list[Any]] = None
+
+
+class AdditionalInformationPayload(msgspec.Struct):
+    """Serialized dictionary payload for additional_information.
+
+    Keys are strings; values are encoded as AdditionalInformationEntry.
+    """
+
+    entries: dict[str, AdditionalInformationEntry]
+
+
+class OmniEngineCoreRequest(EngineCoreRequest):
+    # Optional prompt embeddings (direct-transfer version)
+    prompt_embeds: Optional[PromptEmbedsPayload] = None
+    # Optional additional information dictionary (serialized)
+    additional_information: Optional[AdditionalInformationPayload] = None

vllm_omni/engine/arg_utils.py

@@ -0,0 +1,65 @@
+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"
+    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
+
+        # Create and return the OmniModelConfig instance
+        omni_config = OmniModelConfig(**config_dict)
+        omni_config.hf_config.architectures = omni_config.architectures
+
+        return omni_config

vllm_omni/engine/output_processor.py

@@ -0,0 +1,420 @@
+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
+from vllm.transformers_utils.tokenizer import AnyTokenizer
+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.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,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.mm_type: Optional[str] = None
+        self.mm_accumulated: Optional[torch.Tensor] = None
+
+    @classmethod
+    def from_new_request(
+        cls,
+        tokenizer: AnyTokenizer,
+        request: EngineCoreRequest,
+        prompt: Optional[str],
+        parent_req: Optional[ParentRequest],
+        request_index: int,
+        queue: Optional[Any],
+        log_stats: bool,
+    ) -> "OmniRequestState":
+        if sampling_params := request.sampling_params:
+            if not sampling_params.detokenize:
+                tokenizer = None
+            output_kind = sampling_params.output_kind
+            logprobs_processor = LogprobsProcessor.from_new_request(
+                tokenizer=tokenizer,
+                request=request,
+            )
+            detokenizer = IncrementalDetokenizer.from_new_request(
+                tokenizer=tokenizer,
+                request=request,
+            )
+            max_tokens_param = sampling_params.max_tokens
+            top_p = sampling_params.top_p
+            n = sampling_params.n
+            temperature = sampling_params.temperature
+        else:
+            logprobs_processor = None
+            detokenizer = None
+            max_tokens_param = None
+            top_p = None
+            n = None
+            temperature = None
+            assert request.pooling_params is not None
+            output_kind = request.pooling_params.output_kind
+
+        return cls(
+            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),
+            output_kind=output_kind,
+            prompt=prompt,
+            prompt_token_ids=request.prompt_token_ids,
+            prompt_embeds=request.prompt_embeds,
+            logprobs_processor=logprobs_processor,
+            detokenizer=detokenizer,
+            max_tokens_param=max_tokens_param,
+            top_p=top_p,
+            n=n,
+            temperature=temperature,
+            arrival_time=request.arrival_time,
+            queue=queue,
+            log_stats=log_stats,
+        )
+
+    def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], mm_type: Optional[str]) -> None:
+        if tensor is None:
+            return
+        try:
+            if mm_type:
+                self.mm_type = (mm_type or "").lower()
+            t = tensor.detach()
+            try:
+                t = t.to("cpu")
+            except Exception:
+                # Best-effort CPU move; keep original device if conversion fails
+                logger.debug("Failed to move multimodal tensor to CPU", exc_info=True)
+            if self.mm_accumulated is None:
+                self.mm_accumulated = t
+            else:
+                self.mm_accumulated = torch.cat([self.mm_accumulated, t], dim=0)
+        except Exception:
+            # Log and continue without crashing the output pipeline
+            logger.exception("Error accumulating multimodal tensor")
+
+    # Override: do not route to pooling-only path; always create completion
+    # outputs, and attach pooling_result into the CompletionOutput.
+    def make_request_output(
+        self,
+        new_token_ids: list[int],
+        pooling_output: Optional[torch.Tensor],
+        finish_reason: Optional[FinishReason],
+        stop_reason: Optional[Union[int, str]],
+        kv_transfer_params: Optional[dict[str, Any]] = None,
+    ) -> Optional[Union[OmniRequestOutput, PoolingRequestOutput]]:
+        finished = finish_reason is not None
+        final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
+
+        if not finished and final_only:
+            return None
+
+        request_id = self.request_id
+        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)
+
+        if self.parent_req is None:
+            outputs = [output]
+        else:
+            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)
+
+    def _new_completion_output(
+        self,
+        token_ids: list[int],
+        finish_reason: Optional[FinishReason],
+        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)
+        try:
+            if self.mm_accumulated is not None:
+                tensor = self.mm_accumulated
+                try:
+                    tensor = tensor.detach().to("cpu")
+                except Exception:
+                    logger.debug(
+                        "Failed to move accumulated multimodal tensor to CPU",
+                        exc_info=True,
+                    )
+                # Attach on the completion output for downstream consumers.
+                if not hasattr(base_output, "multimodal_output"):
+                    setattr(base_output, "multimodal_output", {})
+                setattr(base_output, "multimodal_output", {self.mm_type: tensor})
+        except Exception:
+            logger.exception("Error in _new_completion_output")
+        return base_output
+
+
+class MultimodalOutputProcessor(VLLMOutputProcessor):
+    """Handles multimodal output processing by normalizing EngineCoreOutput
+    before delegating to the base vLLM OutputProcessor.
+
+    Strategy:
+    - Route by EngineCoreOutput.output_type when present
+      ("image", "text+image", "latents", "text").
+    - Fallback to pooling/text heuristics when output_type is absent.
+    - Mutate EngineCoreOutput in-place to ensure vLLM's base processor can
+      produce the correct RequestOutput/PoolingRequestOutput.
+    - Allow custom per-modality handlers via register_handler().
+    """
+
+    def __init__(
+        self,
+        tokenizer: AnyTokenizer,
+        log_stats: bool,
+        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.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:
+        self.output_handlers[modality.lower()] = handler
+
+    def add_request(
+        self,
+        request: EngineCoreRequest,
+        prompt: Optional[str],
+        parent_req: Optional[ParentRequest] = None,
+        request_index: int = 0,
+        queue: Optional[RequestOutputCollector] = None,
+    ) -> None:
+        request_id = request.request_id
+        if request_id in self.request_states:
+            raise ValueError(f"Request id {request_id} already running.")
+
+        req_state = OmniRequestState.from_new_request(
+            tokenizer=self.tokenizer,
+            request=request,
+            prompt=prompt,
+            parent_req=parent_req,
+            request_index=request_index,
+            queue=queue,
+            log_stats=self.log_stats,
+        )
+        self.request_states[request_id] = req_state
+        self.lora_states.add_request(req_state)
+        if parent_req:
+            self.parent_requests[parent_req.request_id] = parent_req
+
+    def process_outputs(
+        self,
+        engine_core_outputs: list[EngineCoreOutput],
+        engine_core_timestamp: Optional[float] = None,
+        iteration_stats: Optional[IterationStats] = None,
+    ) -> OutputProcessorOutput:
+        self._reqid_to_mm_type.clear()
+        for eco in engine_core_outputs:
+            mm_type = (self.engine_core_output_type or "").lower()
+            if mm_type:
+                self._reqid_to_mm_type[eco.request_id] = mm_type
+            self._route_and_normalize(eco)
+
+        # Build RequestOutputs without delegating to base, so we can keep ids
+        request_outputs: list[Any] = []
+        reqs_to_abort: list[str] = []
+        for eco in engine_core_outputs:
+            req_id = eco.request_id
+            req_state = self.request_states.get(req_id)
+            if req_state is None:
+                continue
+
+            # 1) 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
+            finish_reason = eco.finish_reason
+            stop_reason = eco.stop_reason
+            kv_transfer_params = eco.kv_transfer_params
+            req_state.num_cached_tokens = eco.num_cached_tokens
+            req_state.is_prefilling = False
+
+            # 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)
+            if stop_string:
+                finish_reason = FinishReason.STOP
+                stop_reason = stop_string
+            req_state.logprobs_processor.update_from_output(eco)
+
+            # 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):
+                    req_state.add_multimodal_tensor(pooling_output, mm_type)
+            except Exception:
+                logger.debug(
+                    "Failed to accumulate multimodal tensor for request %s",
+                    req_id,
+                    exc_info=True,
+                )
+
+            # 3) Create RequestOutput objects, forcing combined mode to keep ids
+            pooling_for_make = pooling_output
+            if pooling_output is not None and new_token_ids:
+                # Do not consume pooling path now; keep ids and attach mm later
+                pooling_for_make = None
+
+            ro = req_state.make_request_output(
+                new_token_ids,
+                pooling_for_make,
+                finish_reason,
+                stop_reason,
+                kv_transfer_params,
+            )
+            if ro:
+                # Attach accumulated multimodal payload if any
+                try:
+                    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")
+                if req_state.queue is not None:
+                    req_state.queue.put(ro)
+                else:
+                    request_outputs.append(ro)
+
+            # 4) Free completed
+            if finish_reason is not None:
+                self.request_states.pop(req_id)
+                parent_req = req_state.parent_req
+                if parent_req and not parent_req.child_requests:
+                    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)
+                if self.tracer:
+                    self.do_tracing(eco, req_state, iteration_stats)
+                # Cleanup per-request mm state
+                if isinstance(req_state, OmniRequestState):
+                    req_state.mm_accumulated = None
+                    req_state.mm_type = None
+
+        return OutputProcessorOutput(
+            request_outputs=request_outputs,
+            reqs_to_abort=reqs_to_abort,
+        )
+
+    # ---- routing helpers ----
+    def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
+        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:
+            try:
+                self.output_handlers[output_type](eco)
+                # Fall through to default fixups in case the handler left gaps
+            except Exception:
+                logger.exception("Error in custom output handler for %s", output_type)
+
+        if output_type == "image":
+            self._process_image_output(eco)
+        elif output_type in ("text+image", "text,image", "image+text"):
+            self._process_text_image_output(eco)
+        elif output_type in ("latents", "latent"):
+            self._process_latents_output(eco)
+        elif output_type in ("audio", "speech"):
+            self._process_audio_output(eco)
+        elif output_type == "text":
+            self._process_text_output(eco)
+        else:
+            # Fallback heuristic
+            if eco.pooling_output is not None:
+                self._process_pooling_output(eco)
+            else:
+                self._process_text_output(eco)
+
+    # ---- modality processors ----
+    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"))
+            if tensor is not None:
+                eco.pooling_output = tensor
+
+    def _process_text_image_output(self, eco: EngineCoreOutput) -> None:
+        """Allow text+image outputs. Text path stays as new_token_ids;
+        image/latents route via pooling_output."""
+        # Preserve text tokens as-is; ensure pooling_output carries image/latents
+        if eco.pooling_output is None:
+            tensor = self._extract_from_multimodal_outputs(
+                eco,
+                keys=(
+                    "image",
+                    "images",
+                    "pixel_values",
+                    "pixels",
+                    "latent",
+                    "latents",
+                    "z",
+                ),
+            )
+            if tensor is not None:
+                eco.pooling_output = tensor
+
+    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"))
+            if tensor is not None:
+                eco.pooling_output = tensor
+
+    def _process_audio_output(self, eco: EngineCoreOutput) -> None:
+        """Ensure audio tensors are surfaced via pooling_output."""
+        if eco.pooling_output is None:
+            tensor = self._extract_from_multimodal_outputs(
+                eco, keys=("audio", "audios", "wav", "waveform", "audio_pcm", "pcm")
+            )
+            if tensor is not None:
+                eco.pooling_output = tensor
+
+    def _process_text_output(self, eco: EngineCoreOutput) -> None:
+        """No-op; base processor will detokenize new_token_ids → text."""
+        return
+
+    def _process_pooling_output(self, eco: EngineCoreOutput) -> None:
+        """Optional sanity checks for pooling tensor."""
+        if eco.pooling_output is None:
+            return
+        if not isinstance(eco.pooling_output, torch.Tensor):
+            # Best-effort: convert to tensor if it's a list/ndarray-like
+            try:
+                eco.pooling_output = torch.as_tensor(eco.pooling_output)
+            except Exception:
+                pass
+
+    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
+        for k in keys:
+            v = mm.get(k)
+            if isinstance(v, torch.Tensor):
+                return v
+        # Try the first tensor in the dict as a fallback
+        for v in mm.values():
+            if isinstance(v, torch.Tensor):
+                return v
+        return None

vllm_omni/engine/processor.py

@@ -0,0 +1,235 @@
+import time
+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
+from vllm.multimodal.utils import argsort_mm_positions
+from vllm.platforms import current_platform
+from vllm.pooling_params import PoolingParams
+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,
+    OmniEngineCoreRequest,
+    PromptEmbedsPayload,
+)
+from vllm_omni.inputs.preprocess import OmniInputPreprocessor
+
+logger = init_logger(__name__)
+
+
+class OmniProcessor(Processor):
+    @staticmethod
+    def _dtype_to_name(dtype: torch.dtype) -> str:
+        mapping = {
+            torch.float32: "float32",
+            torch.float: "float32",
+            torch.float16: "float16",
+            torch.half: "float16",
+            torch.bfloat16: "bfloat16",
+            torch.float64: "float64",
+            torch.double: "float64",
+            torch.int64: "int64",
+            torch.long: "int64",
+            torch.int32: "int32",
+            torch.int: "int32",
+            torch.int16: "int16",
+            torch.short: "int16",
+            torch.int8: "int8",
+            torch.uint8: "uint8",
+            torch.bool: "bool",
+        }
+        return mapping.get(dtype, str(dtype).replace("torch.", ""))
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        tokenizer: AnyTokenizer,
+        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
+    ):
+        super().__init__(vllm_config, tokenizer, mm_registry)
+        self.input_preprocessor = OmniInputPreprocessor(
+            self.model_config,
+            self.tokenizer,
+            mm_registry,
+            mm_processor_cache=self.mm_processor_cache,
+        )
+
+    def process_inputs(
+        self,
+        request_id: str,
+        prompt: PromptType,
+        params: Union[SamplingParams, PoolingParams],
+        arrival_time: Optional[float] = None,
+        lora_request: Optional[LoRARequest] = None,
+        tokenization_kwargs: Optional[dict[str, Any]] = None,
+        trace_headers: Optional[Mapping[str, str]] = None,
+        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} is out of range [0, {data_parallel_size}).")
+
+        if arrival_time is None:
+            arrival_time = time.time()
+
+        # Optionally generate multimodal hash overrides to avoid hashing
+        # multimodal data items by their content as their identifiers.
+
+        # NOTE: when users explicitly turn off BOTH prefix caching and input
+        # processing caching, no multimodal features or embeddings will be
+        # reused across requests, therefore identifying multimodal data items
+        # by their content is no longer necessary, and we create uuids with
+        # request id-modality-index as multimodal hash overrides.
+        if (
+            self.model_config.multimodal_config
+            and self.model_config.multimodal_config.mm_processor_cache_gb == 0
+            and not self.cache_config.enable_prefix_caching
+        ):
+            mm_uuids = self._maybe_build_mm_uuids(request_id, prompt)
+        else:
+            # Otherwise, use user-provided uuids as multimodal hash overrides
+            # if provided.
+            self._validate_multi_modal_uuids(prompt)
+            if isinstance(prompt, dict):
+                mm_uuids = prompt.get("multi_modal_uuids")
+            else:
+                mm_uuids = None
+
+        # Process inputs, which includes:
+        # 1. Tokenize text prompt, with LoRA request if one exists.
+        # 2. For multimodal models with a merged preprocessor, preprocess
+        #   multimodal data and expand prompt token ids accordingly.
+        processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
+            prompt,
+            tokenization_kwargs=tokenization_kwargs,
+            mm_uuids=mm_uuids,
+        )
+        current_platform.validate_request(
+            prompt=prompt,
+            params=params,
+            processed_inputs=processed_inputs,
+        )
+        eos_token_id = self.input_preprocessor.get_eos_token_id()
+
+        encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)
+        self._validate_model_inputs(encoder_inputs, decoder_inputs)
+
+        # Mypy does not always properly infer the types of some elements of
+        # 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
+
+        sampling_params = None
+        pooling_params = None
+        if isinstance(params, SamplingParams):
+            # TODO: can we avoid cloning here in multiproc case?
+            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)
+                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)
+            if self.tokenizer is not None:
+                sampling_params.update_from_tokenizer(self.tokenizer)
+        else:
+            pooling_params = params.clone()
+
+        # Multimodal related.
+        mm_features: Optional[list[MultiModalFeatureSpec]] = None
+
+        if decoder_inputs["type"] == "multimodal":
+            decoder_mm_inputs = decoder_inputs["mm_kwargs"]
+            decoder_mm_positions = decoder_inputs["mm_placeholders"]
+            decoder_mm_hashes = decoder_inputs["mm_hashes"]
+
+            # Merge and flatten multimodal placeholders, hashes and inputs
+            # from dictionaries to lists, and sort them by each item's position
+            # in the input sequence.
+            sorted_mm_idxs = argsort_mm_positions(decoder_mm_positions)
+
+            mm_features = []
+            for modality, idx in sorted_mm_idxs:
+                mm_features.append(
+                    MultiModalFeatureSpec(
+                        data=decoder_mm_inputs[modality][idx],
+                        modality=modality,
+                        identifier=decoder_mm_hashes[modality][idx],
+                        mm_position=decoder_mm_positions[modality][idx],
+                    )
+                )
+
+        # Serialize prompt_embeds and additional_information if provided
+        # (direct-transfer path)
+        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)")
+            # Move to CPU and ensure contiguous memory for stable serialization
+            pe_cpu = pe.detach().to("cpu").contiguous()
+            seq_len, hidden_size = pe_cpu.shape
+            dtype_str = self._dtype_to_name(pe_cpu.dtype)
+            data_bytes = pe_cpu.numpy().tobytes()
+            prompt_embeds_payload = PromptEmbedsPayload(
+                data=data_bytes,
+                shape=[int(seq_len), int(hidden_size)],
+                dtype=dtype_str,
+            )
+        if "additional_information" in decoder_inputs:  # type: ignore[operator]
+            raw_info: dict[str, Any] = decoder_inputs["additional_information"]  # type: ignore[index]  # noqa: E501
+            entries: dict[str, AdditionalInformationEntry] = {}
+            for key, value in raw_info.items():
+                if isinstance(value, torch.Tensor):
+                    v_cpu = value.detach().to("cpu").contiguous()
+                    dtype_str = self._dtype_to_name(v_cpu.dtype)
+                    data_bytes = v_cpu.numpy().tobytes()
+                    entry = AdditionalInformationEntry(
+                        tensor_data=data_bytes,
+                        tensor_shape=[int(x) for x in list(v_cpu.shape)],
+                        tensor_dtype=dtype_str,
+                    )
+                elif isinstance(value, list):
+                    entry = AdditionalInformationEntry(list_data=value)
+                else:
+                    raise ValueError("additional_information values must be Tensor or list")
+                entries[key] = entry
+            additional_information_payload = AdditionalInformationPayload(entries=entries)
+
+        return prompt_str, OmniEngineCoreRequest(
+            request_id=request_id,
+            prompt_token_ids=prompt_token_ids,
+            mm_features=mm_features,
+            sampling_params=sampling_params,
+            pooling_params=pooling_params,
+            eos_token_id=eos_token_id,
+            arrival_time=arrival_time,
+            lora_request=lora_request,
+            cache_salt=decoder_inputs.get("cache_salt"),
+            priority=priority,
+            data_parallel_rank=data_parallel_rank,
+            trace_headers=trace_headers,
+            prompt_embeds=prompt_embeds_payload,
+            additional_information=additional_information_payload,
+        )

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/__init__.py

@@ -0,0 +1,5 @@
+"""CLI helpers for vLLM-omni entrypoints."""
+
+from .serve import OmniServeCommand
+
+__all__ = ["OmniServeCommand"]

vllm_omni/entrypoints/cli/main.py

@@ -0,0 +1,57 @@
+"""
+CLI entry point for vLLM-omni that intercepts vLLM commands.
+"""
+
+import importlib.metadata
+import sys
+
+
+def main():
+    """Main CLI entry point that intercepts vLLM commands."""
+    # Check if --omni flag is present
+    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__":
+    main()

vllm_omni/entrypoints/cli/serve.py

@@ -0,0 +1,67 @@
+"""
+Omni serve command for vLLM-omni.
+"""
+
+import argparse
+
+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(CLISubcommand):
+    """The `serve` subcommand for the vLLM CLI."""
+
+    name = "serve"
+
+    @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]",
+        )
+
+        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="Enable vLLM-omni mode",
+        )
+        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
+
+
+def cmd_init() -> list[CLISubcommand]:
+    return [OmniServeCommand()]

vllm_omni/entrypoints/log_utils.py

@@ -0,0 +1,693 @@
+from __future__ import annotations
+
+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: str | None, num_stages: int) -> None:
+    try:
+        if not log_file:
+            return
+        try:
+            if os.path.exists(log_file):
+                os.remove(log_file)
+        except Exception:
+            pass
+        # Per-stage logs and stats
+        for sid in range(num_stages):
+            try:
+                p = f"{log_file}.stage{sid}.log"
+                if os.path.exists(p):
+                    os.remove(p)
+            except Exception:
+                pass
+            try:
+                p = f"{log_file}.stage{sid}.stats.jsonl"
+                if os.path.exists(p):
+                    os.remove(p)
+            except Exception:
+                pass
+        # Orchestrator stats files
+        try:
+            p = f"{log_file}.orchestrator.stats.jsonl"
+            if os.path.exists(p):
+                os.remove(p)
+        except Exception:
+            pass
+        try:
+            p = f"{log_file}.overall.stats.jsonl"
+            if os.path.exists(p):
+                os.remove(p)
+        except Exception:
+            pass
+    except Exception:
+        pass
+
+
+def configure_orchestrator_logger(logger: logging.Logger, log_file: str | None) -> None:
+    try:
+        if not log_file:
+            return
+        has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers)
+        if not has_file_handler:
+            fh = logging.FileHandler(log_file)
+            fh.setLevel(logging.DEBUG)
+            fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s"))
+            logger.addHandler(fh)
+            logger.setLevel(logging.DEBUG)
+    except Exception:
+        pass
+
+
+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"
+            overall_file = f"{log_file}.overall.stats.jsonl"
+    except Exception:
+        stats_file = None
+        overall_file = None
+    return stats_file, overall_file
+
+
+def _safe_append_jsonl(path: str | None, record: dict[str, Any]) -> None:
+    if not path:
+        return
+    try:
+        _append_jsonl(path, record)  # type: ignore[arg-type]
+    except Exception:
+        pass
+
+
+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: 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: 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: 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: str | None, summary: dict[str, Any]) -> None:
+    _safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary})
+
+
+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: str | None, record: dict[str, Any]) -> None:
+    _safe_append_jsonl(overall_stats_file, record)
+
+
+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: 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: 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: 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
+    )
+    log_stage_request_stats(
+        stats_file,
+        stage_id,
+        request_id,
+        int(batch_size),
+        int(num_tokens),
+        float(stage_gen_time_ms),
+        float(tokens_per_s),
+        int(rx_transfer_bytes),
+        float(rx_decode_time_ms),
+        float(rx_mbps),
+    )
+
+
+# ----------------- 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:
+    try:
+        stage_req_counts[stage_id] += 1
+        stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
+        rid_int = int(req_id)
+        pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+        pr_stages = pr["stages"]  # type: ignore[index]
+        pr_stages[stage_id] = {
+            "stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)),
+            "num_tokens_out": int(metrics.get("num_tokens_out", 0)),
+        }
+    except Exception:
+        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,
+) -> 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,
+                }
+                transfer_agg[key] = agg
+            agg["sum_rx_bytes"] += float(rx_bytes)
+            agg["sum_rx_ms"] += float(rx_ms)
+            agg["rx_count"] += 1.0
+
+            # Try combine with sender-side timing if present
+            try:
+                rid_int = int(req_id)
+            except Exception:
+                rid_int = req_id  # best effort
+            s = transfer_edge_req.get((stage_id - 1, stage_id, rid_int))
+            if s is None:
+                return None
+            tx_ms = float(s.get("tx_ms", 0.0))
+            size_b = float(s.get("size_bytes", rx_bytes))
+            total_ms = tx_ms + float(in_flight_ms) + float(rx_ms)
+            agg["sum_total_ms"] += total_ms
+            agg["total_count"] += 1.0
+            # accumulate per-request transfer totals
+            try:
+                pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+                pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms  # type: ignore[index]
+                pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes)  # type: ignore[index]
+            except Exception:
+                pass
+            return int(size_b), float(tx_ms), float(total_ms)
+        return None
+    except Exception:
+        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:
+    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,
+            }
+            transfer_agg[key] = agg
+        agg["sum_bytes"] += float(size_bytes)
+        agg["sum_ms"] += float(tx_ms)
+        agg["count"] += 1.0
+        # Store sender-side timing for per-request combination
+        try:
+            rid_int = int(req_id)
+        except Exception:
+            rid_int = req_id
+        transfer_edge_req[(from_stage, to_stage, rid_int)] = {
+            "tx_ms": float(tx_ms),
+            "size_bytes": float(size_bytes),
+        }
+    except Exception:
+        pass
+
+
+def count_tokens_from_outputs(engine_outputs: list[Any]) -> int:
+    total = 0
+    for _ro in engine_outputs:
+        try:
+            outs = getattr(_ro, "outputs", None)
+            if outs and len(outs) > 0:
+                tokens = getattr(outs[0], "token_ids", None)
+                if tokens is not None:
+                    total += len(tokens)
+        except Exception:
+            pass
+    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]] = []
+    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,
+            }
+        )
+    return summary
+
+
+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))
+        samples = int(agg.get("count", 0.0))
+        tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
+        sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0))
+        sum_rx_ms = float(agg.get("sum_rx_ms", 0.0))
+        samples_rx = int(agg.get("rx_count", 0.0))
+        rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0
+        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,
+            }
+        )
+    return summary
+
+
+class OrchestratorMetrics:
+    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.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.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[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:
+        record_stage_metrics(
+            self.per_request,
+            self.stage_req_counts,
+            self.stage_total_time_ms,
+            self.stage_total_tokens,
+            stage_id,
+            req_id,
+            metrics,
+        )
+        try:
+            batch_id_raw = metrics.get("batch_id", None)
+            if batch_id_raw is not None:
+                batch_id = int(batch_id_raw)
+                if batch_id not in self.stage_seen_batches[stage_id]:
+                    self.stage_total_time_ms[stage_id] += float(metrics.get("stage_gen_time_ms", 0.0))
+                    self.stage_seen_batches[stage_id].add(batch_id)
+        except Exception:
+            pass
+        rx_b = float(metrics.get("rx_transfer_bytes", 0.0))
+        rx_ms = float(metrics.get("rx_decode_time_ms", 0.0))
+        in_flight_ms = float(metrics.get("rx_in_flight_time_ms", 0.0))
+        combined = aggregate_rx_and_maybe_total(
+            self.transfer_edge_req,
+            self.transfer_agg,
+            self.per_request,
+            stage_id,
+            req_id,
+            rx_b,
+            rx_ms,
+            in_flight_ms,
+        )
+        if self.enable_stats and self.stats_file and stage_id > 0:
+            log_transfer_rx(
+                self.stats_file,
+                stage_id - 1,
+                stage_id,
+                req_id,
+                int(rx_b),
+                float(rx_ms),
+                float(in_flight_ms),
+            )
+            if combined is not None:
+                size_b_c, tx_ms_c, total_ms_c = combined
+                log_transfer_total(
+                    self.stats_file,
+                    stage_id - 1,
+                    stage_id,
+                    req_id,
+                    int(size_b_c),
+                    float(tx_ms_c),
+                    float(in_flight_ms),
+                    float(rx_ms),
+                    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:
+        # 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()
+        if self.enable_stats and self.stats_file:
+            log_transfer_tx(
+                self.stats_file,
+                from_stage,
+                to_stage,
+                req_id,
+                int(size_bytes),
+                float(tx_ms),
+                bool(used_shm),
+            )
+        record_sender_transfer_agg(
+            self.transfer_agg,
+            self.transfer_edge_req,
+            from_stage,
+            to_stage,
+            req_id,
+            int(size_bytes),
+            float(tx_ms),
+        )
+
+    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:
+            rid_int = req_id
+        _t0 = float(req_start_ts)
+        _t1 = time.time()
+        # Update last output time for this stage
+        prev_last = self.stage_last_ts[stage_id]
+        self.stage_last_ts[stage_id] = _t1 if prev_last is None else max(prev_last, _t1)
+        self.last_finish_ts = max(self.last_finish_ts, _t1)
+        e2e_ms = (_t1 - _t0) * 1000.0
+        num_tokens = count_tokens_from_outputs(engine_outputs)
+        self.e2e_total_ms += e2e_ms
+        self.e2e_total_tokens += int(num_tokens)
+        self.e2e_count += 1
+        self.e2e_done.add(rid_int)
+        pr = self.per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+        per_req_record = {
+            "type": "overall_request",
+            "request_id": rid_int,
+            "e2e_time_ms": e2e_ms,
+            "num_tokens_out": int(num_tokens),
+            "transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)),
+            "transfers_total_bytes": int(pr.get("transfers_bytes", 0)),
+            "stages": pr.get("stages", {}),
+        }
+        self.sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0))
+        if self.enable_stats and self.overall_stats_file:
+            log_overall_record(self.overall_stats_file, per_req_record)
+        if self.enable_stats and self.stats_file:
+            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]:
+        # Compute stage summary using wall time between first input and last output per stage
+        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
+            )
+            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),
+                }
+            )
+        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] = {
+            "e2e_requests": int(self.e2e_count),
+            "e2e_total_time_ms": float(wall_time_ms),
+            "e2e_sum_time_ms": float(self.e2e_total_ms),
+            "e2e_total_tokens": int(self.e2e_total_tokens),
+            "e2e_avg_time_per_request_ms": e2e_avg_req,
+            "e2e_avg_tokens_per_s": e2e_avg_tok,
+            "wall_time_ms": wall_time_ms,
+            "final_stage_id": final_stage_id_for_e2e,
+            "stages": stage_summary,
+            "transfers": transfer_summary,
+        }
+        if self.enable_stats and self.stats_file:
+            log_orchestrator_summary(self.stats_file, summary)
+        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

@@ -0,0 +1,559 @@
+import multiprocessing as mp
+import os
+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
+
+# External library imports (vLLM)
+from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field
+from vllm.engine.arg_utils import HfOverrides
+from vllm.entrypoints.llm import LLM
+from vllm.inputs import PromptType
+from vllm.logger import init_logger
+from vllm.plugins.io_processors import get_io_processor
+from vllm.sampling_params import SamplingParams
+from vllm.usage.usage_lib import UsageContext
+from vllm.utils import Counter
+from vllm.v1.engine.llm_engine import LLMEngine
+
+# Internal imports (our code)
+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.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 OmniLLM:
+    """Multi-process pipelined OmniLLM.
+
+    - Per-stage process with copy-based IPC (Queues)
+    - window=-1 across stages (downstream starts when upstream finishes)
+    - max_inflight=1 per stage (serial within a stage), but pipeline across stages
+    """
+
+    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_path is None:
+            self.stage_configs = load_stage_configs_from_model(model)
+        else:
+            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
+
+        # 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]
+        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,
+                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 generate(
+        self,
+        prompts: Union[PromptType, Sequence[PromptType]],
+        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)}")
+
+        # Normalize prompts to a list for per-request iteration
+        if not isinstance(prompts, (list, tuple)):
+            request_prompts: list[PromptType] = [prompts]
+        else:
+            request_prompts = list(prompts)
+
+        final_outputs: list[OmniRequestOutput] = []
+
+        # Orchestrator keeps stage objects for input derivation
+        num_stages = len(self.stage_list)
+
+        # Map from request_id to original prompt
+        request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
+
+        # Track per-request start time for end-to-end timing
+        _req_start_ts: dict[int, float] = {}
+        _wall_start_ts: float = time.time()
+
+        # 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 %d requests into stage-0", len(request_prompts))
+        # Mark first input time for stage-0
+        metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
+
+        for req_id, prompt in request_id_to_prompt.items():
+            sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
+            task = {
+                "request_id": req_id,
+                "engine_inputs": prompt,
+                "sampling_params": sp0,
+            }
+            self.stage_list[0].submit(task)
+            _req_start_ts[req_id] = time.time()
+            logger.debug("[Orchestrator] Enqueued request %s to stage-0", req_id)
+
+        # For each stage, forward results to next stage; collect finals at the end
+        # We pipeline by continually polling output queues in stage order
+        remaining_by_stage: list[int] = [len(request_prompts)] + [0] * (num_stages - 1)
+        completed_requests = 0
+        total_requests = len(request_prompts)
+
+        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):
+                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):
+                    final_outputs.append(
+                        OmniRequestOutput(
+                            stage_id=stage_id,
+                            final_output_type=stage.final_output_type,  # type: ignore[attr-defined]
+                            request_output=engine_outputs,
+                        )
+                    )
+                    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,
+                        )
+
+                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, [request_id_to_prompt[req_id]])
+                    sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
+                    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,
+                    )
+                    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,
+                    )
+
+            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)
+
+        return final_outputs
+
+    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)
+                    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.",
+                    "Check model weights path and network reachability (if loading remotely).",
+                    "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)
+
+
+class OmniStageLLM(LLM):
+    def __init__(
+        self,
+        model: str,
+        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,
+        **kwargs,
+    ):
+        """LLM constructor."""
+        if "disable_log_stats" not in kwargs:
+            kwargs["disable_log_stats"] = True
+
+        if "worker_cls" in kwargs:
+            worker_cls = kwargs["worker_cls"]
+            # if the worker_cls is not qualified string name,
+            # we serialize it using cloudpickle to avoid pickling issues
+            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):
+            from vllm.config.kv_transfer import KVTransferConfig
+
+            raw_config_dict = kwargs["kv_transfer_config"]
+            try:
+                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",
+                    raw_config_dict,
+                    e,
+                )
+                # Consider re-raising a more specific vLLM error or ValueError
+                # to provide better context to the user.
+                raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e
+
+        if hf_overrides is None:
+            hf_overrides = {}
+
+        if compilation_config is not None:
+            if isinstance(compilation_config, int):
+                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)}
+                )
+            else:
+                compilation_config_instance = compilation_config
+        else:
+            compilation_config_instance = CompilationConfig()
+
+        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)}
+                )
+            else:
+                structured_outputs_instance = structured_outputs_config
+        else:
+            structured_outputs_instance = StructuredOutputsConfig()
+
+        engine_args = OmniEngineArgs(
+            model=model,
+            hf_overrides=hf_overrides,
+            compilation_config=compilation_config_instance,
+            structured_outputs_config=structured_outputs_instance,
+            **kwargs,
+        )
+
+        # 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.output_processor = MultimodalOutputProcessor(
+            tokenizer=self.llm_engine.tokenizer,
+            log_stats=self.llm_engine.log_stats,
+            engine_core_output_type=engine_args.engine_output_type,
+        )
+        self.llm_engine.processor = OmniProcessor(
+            vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer
+        )
+        self.engine_class = type(self.llm_engine)
+
+        self.request_counter = Counter()
+        self.default_sampling_params: Union[dict[str, Any], None] = None
+
+        supported_tasks = self.llm_engine.get_supported_tasks()  # type: ignore
+
+        logger.info("Supported_tasks: %s", supported_tasks)
+
+        self.supported_tasks = supported_tasks
+
+        # 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)

vllm_omni/entrypoints/omni_stage.py

@@ -0,0 +1,723 @@
+"""
+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, 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
+        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 = importlib.import_module(module_path)
+            self.custom_process_input_func = getattr(module, func_name)
+        else:
+            self.custom_process_input_func = 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
+        self._proc: Optional[mp.Process] = None
+        self._log_file: Optional[str] = None
+        self._shm_threshold_bytes: int = 65536
+        self._logger = logging.getLogger(__name__)
+
+    def set_engine(self, engine: LLMEngine) -> None:
+        """Initialize the engine for the stage."""
+        self.engine = engine
+
+    def set_async_engine(self, async_engine: AsyncLLM) -> None:
+        """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
+
+    # ----------------- New Orchestration APIs -----------------
+    def attach_queues(self, in_q: mp.Queue, out_q: mp.Queue) -> None:
+        self._in_q = in_q
+        self._out_q = out_q
+
+    def init_stage_worker(
+        self,
+        model: str,
+        *,
+        is_async: bool = False,
+        log_file: Optional[str] = None,
+        shm_threshold_bytes: int = 65536,
+        ctx: Optional[mp.context.BaseContext] = None,
+        batch_timeout: int = 10,
+    ) -> None:
+        assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process"
+        self._log_file = log_file
+        self._shm_threshold_bytes = shm_threshold_bytes
+        ctx = ctx or mp.get_context("spawn")
+        # 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_id": self.stage_id,
+            "engine_args": engine_args,
+            "runtime": runtime_cfg,
+            "shm_threshold_bytes": self._shm_threshold_bytes,
+        }
+        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:
+        if self._in_q is not None:
+            try:
+                self._in_q.put_nowait(None)
+            except Exception as e:
+                self._logger.warning("[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e)
+        if self._proc is not None:
+            try:
+                self._proc.join(timeout=5)
+            except Exception as e:
+                self._logger.debug("[Stage-%s] join() failed: %s", self.stage_id, e, exc_info=True)
+            if self._proc.is_alive():
+                try:
+                    self._proc.terminate()
+                except Exception as e:
+                    self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
+
+    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]]:
+        assert self._out_q is not None
+        try:
+            return self._out_q.get_nowait()
+        except Exception:
+            return None
+
+    def process_engine_inputs(
+        self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None
+    ) -> list[Union[OmniTokensPrompt, TextPrompt]]:
+        """Process the engine input for the stage."""
+        if self.custom_process_input_func is None:
+            engine_inputs = []
+            if len(self.engine_input_source) == 0:
+                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: 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),
+                )
+                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)
+
+
+def _stage_worker(
+    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.log_utils import (  # noqa: WPS433
+        compute_and_log_stage_request_stats,
+        count_tokens_from_outputs,
+        log_stage_batch_stats,
+        log_stage_running_avg,
+    )
+    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"]
+    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")
+            )
+
+            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()),
+    )
+    stage_engine = OmniStageLLM(model=model, **engine_args)
+    _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
+    # Signal readiness to orchestrator
+    try:
+        out_q.put({"type": "stage_ready", "stage_id": stage_id})
+    except Exception:
+        pass
+
+    # 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
+
+        max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
+        batch_tasks: list[dict[str, Any]] = [task]
+        if max_batch_size > 1:
+            while len(batch_tasks) < max_batch_size:
+                if not in_q.empty():
+                    extra = in_q.get(timeout=batch_timeout)
+                    if extra is None:
+                        in_q.put(None)
+                        break
+                    batch_tasks.append(extra)
+                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] = {}
+        for t in batch_tasks:
+            rid = t["request_id"]
+            try:
+                sent_ts = float(t.get("sent_ts", None)) if isinstance(t, 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(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)
+            if isinstance(ein, list):
+                batch_engine_inputs.extend(ein)
+            elif isinstance(ein, dict):
+                batch_engine_inputs.append(ein)
+            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,
+        )
+        print("--------------------------------", 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_t0 = _time.time()
+            for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False):
+                gen_outputs.append(ro)
+            _gen_t1 = _time.time()
+            _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] = []
+            for ro in gen_outputs:
+                rid = getattr(ro, "request_id", None)
+                if rid in req_to_outputs:
+                    req_to_outputs[rid].append(ro)
+                else:
+                    unmapped.append(ro)
+            if unmapped:
+                idx = 0
+                for ro in unmapped:
+                    target_rid = batch_request_ids[idx % len(batch_request_ids)]
+                    req_to_outputs[target_rid].append(ro)
+                    idx += 1
+
+            # Per-request stats logging and aggregates
+            for rid in batch_request_ids:
+                _r_outputs = req_to_outputs.get(rid, [])
+                _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,
+                    len(batch_tasks),
+                    float(_gen_ms),
+                    list(batch_request_ids),
+                )
+
+            # Emit per-request results
+            for rid in batch_request_ids:
+                r_outputs = req_to_outputs.get(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,
+                            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]
+                        )
+                    if use_shm:
+                        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(
+                        {
+                            "request_id": rid,
+                            "stage_id": stage_id,
+                            "engine_outputs": payload,
+                            "metrics": _metrics,
+                        }
+                    )
+            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,
+                        "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 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

@@ -0,0 +1,271 @@
+from __future__ import annotations
+
+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: str | int | None) -> None:
+    """Configure per-stage CUDA visibility and current device.
+
+    Behavior
+    - Comma-separated string (e.g. "2,5,7"): set CUDA_VISIBLE_DEVICES exactly
+      to this list; logical index 0 is used as current device.
+    - Integer or digit-string: treat as logical index (0-based) into the current
+      CUDA_VISIBLE_DEVICES mapping; map to the physical device, and then set
+      CUDA_VISIBLE_DEVICES to this single device.
+    - None/"cpu": keep default visibility.
+    - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string.
+    """
+    try:
+        selected_physical: int | None = None
+        logical_idx: int | None = None
+
+        if isinstance(devices, str) and "," in devices:
+            os.environ["CUDA_VISIBLE_DEVICES"] = devices
+            toks = [t.strip() for t in devices.split(",") if t.strip() != ""]
+            if toks:
+                try:
+                    selected_physical = int(toks[0])
+                    logger.debug(
+                        "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s",
+                        stage_id,
+                        devices,
+                        selected_physical,
+                    )
+                except Exception as e:
+                    logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
+                    selected_physical = None
+        elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()):
+            logical_idx = max(0, int(devices))
+            vis = os.environ.get("CUDA_VISIBLE_DEVICES")
+            if vis:
+                try:
+                    mapping = [int(x) for x in vis.split(",") if x.strip() != ""]
+                    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,
+                    )
+                    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,
+            )
+        elif devices in (None, "cpu"):
+            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,
+            )
+
+        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,
+                    )
+                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),
+                        }
+                    )
+                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,
+            )
+    except Exception as e:
+        logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)
+
+
+def serialize_obj(obj: Any) -> bytes:
+    """Serialize a Python object to bytes using cloudpickle."""
+    return cloudpickle.dumps(obj)
+
+
+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.
+    """
+    shm = _shm.SharedMemory(create=True, size=len(payload))
+    mv = memoryview(shm.buf)
+    mv[: len(payload)] = payload
+    del mv
+    meta = {"name": shm.name, "size": len(payload)}
+    try:
+        shm.close()
+    except Exception as e:
+        logger.debug("Failed to close shared memory: %s", e)
+    return meta
+
+
+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)
+    data = bytes(mv[: meta["size"]])
+    del mv
+    try:
+        shm.close()
+    except Exception:
+        pass
+    try:
+        shm.unlink()
+    except Exception:
+        pass
+    return data
+
+
+def _ensure_parent_dir(path: str) -> None:
+    """Ensure the parent directory for a file path exists (best-effort)."""
+    try:
+        parent = os.path.dirname(path)
+        if parent:
+            os.makedirs(parent, exist_ok=True)
+    except Exception:
+        pass
+
+
+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
+    to a distinct file. For concurrent writes to the same file, OS append
+    semantics typically suffice, but no additional locking is provided.
+    """
+    try:
+        _ensure_parent_dir(path)
+        line = json.dumps(record, ensure_ascii=False)
+        fd = os.open(path, os.O_APPEND | os.O_CREAT | os.O_WRONLY, 0o644)
+        with os.fdopen(fd, "a", encoding="utf-8") as f:
+            f.write(line + "\n")
+    except Exception:
+        logger.exception("Failed to append JSONL to %s", path)
+
+
+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).
+    """
+    payload = serialize_obj(obj)
+    if len(payload) > threshold:
+        return True, shm_write_bytes(payload)
+    return False, obj
+
+
+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
+    decode-time and size metrics.
+    """
+    if shm_key in container:
+        return pickle.loads(shm_read_bytes(container[shm_key]))
+    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]]:
+    """Load object and return (object, metrics) with RX bytes and decode time.
+
+    Metrics keys:
+      - rx_transfer_bytes: int
+      - 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]
+        payload = shm_read_bytes(meta)
+        obj = pickle.loads(payload)
+        try:
+            rx_bytes = int(meta.get("size", len(payload)))  # type: ignore[call-arg]
+        except Exception:
+            rx_bytes = len(payload)
+    else:
+        obj = container[obj_key]
+        try:
+            rx_bytes = len(serialize_obj(obj))
+        except Exception:
+            rx_bytes = 0
+    t1 = _time.time()
+    rx_decode_ms = (t1 - t0) * 1000.0
+    return obj, {
+        "rx_transfer_bytes": int(rx_bytes),
+        "rx_decode_time_ms": float(rx_decode_ms),
+    }
+
+
+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] = {}
+    use_shm, data = maybe_dump_to_shm(obj, threshold)
+    if use_shm:
+        payload[shm_key] = data
+    else:
+        payload[obj_key] = data
+    return payload
+
+
+# Convert OmegaConf/objects to plain dicts
+def _to_dict(x: Any) -> dict[str, Any]:
+    try:
+        if isinstance(x, dict):
+            return dict(x)
+        return OmegaConf.to_container(x, resolve=True)  # type: ignore[arg-type]
+    except Exception:
+        try:
+            return dict(x)
+        except Exception:
+            return {}

vllm_omni/entrypoints/utils.py

@@ -0,0 +1,31 @@
+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
+
+
+def load_stage_configs_from_model(model: str):
+    """Load stage configs from model."""
+    hf_config = get_config(model, trust_remote_code=True)
+    model_type = hf_config.model_type
+    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
+    stage_config_path = PROJECT_ROOT / stage_config_file
+    if not os.path.exists(stage_config_path):
+        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
+    stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path))
+    return stage_configs
+
+
+def load_stage_configs_from_yaml(config_path: str):
+    """Load stage configs from yaml file."""
+    config_data = OmegaConf.load(config_path)
+    return config_data.stage_args

vllm_omni/inputs/data.py

@@ -0,0 +1,59 @@
+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
+
+
+class OmniTokensPrompt(TokensPrompt):
+    prompt_embeds: NotRequired[torch.Tensor]
+    """The embeddings of the prompt."""
+
+    # New: optional additional information dictionary
+    # Values may be torch.Tensor or list
+    additional_information: NotRequired[dict[str, Any]]
+
+
+class OmniTokenInputs(TokenInputs):
+    # New: optional prompt embeddings aligned with token ids
+    prompt_embeds: NotRequired[torch.Tensor]
+
+    # New: optional additional information dictionary
+    # Values may be torch.Tensor or list
+    additional_information: NotRequired[dict[str, Any]]
+
+
+class OmniEmbedsPrompt(EmbedsPrompt):
+    # New: optional prompt embeddings aligned with token ids
+    prompt_embeds: NotRequired[torch.Tensor]
+
+    # New: optional additional information dictionary
+    # Values may be torch.Tensor or list
+    additional_information: NotRequired[dict[str, Any]]
+
+
+def token_inputs_omni(
+    prompt_token_ids: list[int],
+    prompt: Optional[str] = None,
+    cache_salt: Optional[str] = None,
+    prompt_embeds: Optional[torch.Tensor] = None,
+    additional_information: Optional[dict[str, Any]] = None,
+) -> OmniTokenInputs:
+    """Construct token inputs with optional embeddings and metadata."""
+    inputs = OmniTokenInputs(type="token", prompt_token_ids=prompt_token_ids)
+
+    if prompt is not None:
+        inputs["prompt"] = prompt
+    if cache_salt is not None:
+        inputs["cache_salt"] = cache_salt
+    if prompt_embeds is not None:
+        inputs["prompt_embeds"] = prompt_embeds
+    if additional_information is not None:
+        inputs["additional_information"] = additional_information
+
+    return inputs

vllm_omni/inputs/parse.py

@@ -0,0 +1,25 @@
+from vllm.inputs.data import SingletonPrompt
+from vllm.inputs.parse import (
+    ParsedEmbedsPrompt,
+    ParsedSingletonPrompt,
+    ParsedStrPrompt,
+    ParsedTextPrompt,
+    ParsedTokensPrompt,
+)
+
+
+def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPrompt:
+    if isinstance(prompt, str):
+        return ParsedStrPrompt(type="str", content=prompt)
+    elif isinstance(prompt, dict):
+        # Type ignores are because mypy does not correctly infer the TypedDicts
+        # Pyright does succeed.
+        # 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]
+        elif "prompt_embeds" in prompt:
+            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")

vllm_omni/inputs/preprocess.py

@@ -0,0 +1,90 @@
+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(
+        self,
+        parsed_content: TokensPrompt,
+        tokenization_kwargs: Optional[dict[str, Any]] = None,
+        *,
+        mm_uuids: Optional[MultiModalUUIDDict] = None,
+    ) -> Union[OmniTokenInputs, MultiModalInputs]:
+        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")
+
+        inputs: Union[OmniTokenInputs, MultiModalInputs]
+        if multi_modal_data := parsed_content.get("multi_modal_data"):
+            inputs = self._process_multimodal(
+                prompt_token_ids,
+                multi_modal_data,
+                parsed_content.get("mm_processor_kwargs"),
+                tokenization_kwargs=tokenization_kwargs,
+                mm_uuids=mm_uuids,
+            )
+        else:
+            inputs = token_inputs_omni(
+                prompt_token_ids=prompt_token_ids,
+                prompt_embeds=prompt_embeds,
+                additional_information=additional_information,
+            )
+
+        if cache_salt := parsed_content.get("cache_salt"):
+            inputs["cache_salt"] = cache_salt
+
+        return inputs
+
+    def _prompt_to_llm_inputs(
+        self,
+        prompt: SingletonPrompt,
+        tokenization_kwargs: Optional[dict[str, Any]] = None,
+        *,
+        mm_uuids: Optional[MultiModalUUIDDict] = None,
+    ) -> SingletonInputs:
+        """
+        Extract the singleton inputs from a prompt.
+
+        Arguments:
+
+        * prompt: single encoder or decoder input prompt
+        * lora_request: this is only valid for decoder prompts
+        * return_mm_hashes: whether to return multimodal hashes
+
+        Returns:
+
+        * Input container compatible with vLLM's singleton prompt handling.
+        """
+        parsed = parse_singleton_prompt_omni(prompt)
+
+        if parsed["type"] == "embeds":
+            return self._process_embeds(parsed["content"])
+        if parsed["type"] == "tokens":
+            return self._process_tokens(
+                parsed["content"],
+                mm_uuids=mm_uuids,
+            )
+        if parsed["type"] == "text":
+            return self._process_text(
+                parsed["content"],
+                tokenization_kwargs=tokenization_kwargs,
+                mm_uuids=mm_uuids,
+            )
+        if parsed["type"] == "str":
+            return self._process_text(
+                TextPrompt(prompt=parsed["content"]),
+                tokenization_kwargs=tokenization_kwargs,
+                mm_uuids=mm_uuids,
+            )
+
+        assert_never(parsed)

vllm_omni/model_executor/layers/mrope.py

@@ -0,0 +1,675 @@
+import itertools
+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
+
+
+def _apply_rotary_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    is_neox_style: bool,
+) -> torch.Tensor:
+    """
+    Args:
+        x: [num_tokens, num_heads, head_size]
+        cos: [num_tokens, head_size // 2]
+        sin: [num_tokens, head_size // 2]
+        is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
+            positional embeddings.
+    """
+    cos = cos.unsqueeze(-2).to(x.dtype)
+    sin = sin.unsqueeze(-2).to(x.dtype)
+    if is_neox_style:
+        x1, x2 = torch.chunk(x, 2, dim=-1)
+    else:
+        x1 = x[..., ::2]
+        x2 = x[..., 1::2]
+    o1 = x1 * cos - x2 * sin
+    o2 = x2 * cos + x1 * sin
+    if is_neox_style:
+        return torch.cat((o1, o2), dim=-1)
+    else:
+        return torch.stack((o1, o2), dim=-1).flatten(-2)
+
+
+class MRotaryEmbedding(RotaryEmbedding):
+    """Rotary Embedding with Multimodal Sections."""
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: float,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+        mrope_section: Optional[list[int]] = None,
+        mrope_interleaved: bool = False,
+        # YaRN parameters.
+        *,
+        scaling_factor: Optional[float] = None,
+        extrapolation_factor: float = 1,
+        attn_factor: float = 1,
+        beta_fast: int = 32,
+        beta_slow: int = 1,
+    ) -> None:
+        # In Qwen2.5-VL, the maximum index value is related to the duration of
+        # the input video. We enlarge max_position_embeddings to 4 times to get
+        # a larger the cos and sin cache.
+        self.cache_max_position_num = max_position_embeddings * 4
+        super().__init__(
+            head_size,
+            rotary_dim,
+            self.cache_max_position_num,
+            base,
+            is_neox_style,
+            dtype,
+        )
+
+        self.mrope_section = mrope_section
+        if self.mrope_section:
+            print(
+                "Warning: mrope_section check is disabled in Qwen2.5-Omni, "
+                "this may cause errors, and should be restored in the future."
+            )
+            # assert sum(self.mrope_section) == rotary_dim // 2
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """PyTorch-native implementation equivalent to forward().
+
+        Args:
+            positions:
+                [num_tokens,] (text only) or
+                [3, num_tokens] (T/H/W positions with multimodal inputs)
+            query: [num_tokens, num_heads * head_size]
+            key: [num_tokens, num_kv_heads * head_size]
+        """
+        assert positions.ndim == 1 or positions.ndim == 2
+        assert key is not None
+
+        num_tokens = positions.shape[-1]
+        cos_sin = self.cos_sin_cache[positions]
+        cos, sin = cos_sin.chunk(2, dim=-1)
+        if positions.ndim == 2:
+            assert self.mrope_section
+
+            cos = torch.cat(
+                [m[i] for i, m in enumerate(cos.split(self.mrope_section, dim=-1))],
+                dim=-1,
+            )
+            sin = torch.cat(
+                [m[i] for i, m in enumerate(sin.split(self.mrope_section, dim=-1))],
+                dim=-1,
+            )
+
+        query_shape = query.shape
+        query = query.view(num_tokens, -1, self.head_size)
+        query_rot = query[..., : self.rotary_dim]
+        query_pass = query[..., self.rotary_dim :]
+        # query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin,
+        #                                       self.is_neox_style)
+        query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
+        query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
+
+        key_shape = key.shape
+        key = key.view(num_tokens, -1, self.head_size)
+        key_rot = key[..., : self.rotary_dim]
+        key_pass = key[..., self.rotary_dim :]
+        # key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin,
+        #                                     self.is_neox_style)
+        key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
+        key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
+        return query, key
+
+    @classmethod
+    def get_input_positions(
+        cls,
+        input_tokens: list[int],
+        hf_config: PretrainedConfig,
+        image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
+        video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
+        second_per_grid_ts: Optional[list[float]],
+        context_len: int = 0,
+        seq_len: Optional[int] = None,
+        audio_feature_lengths: Optional[torch.Tensor] = None,
+        use_audio_in_video: bool = False,
+    ) -> tuple[list[list[int]], int]:
+        """Get mrope input positions and delta value."""
+
+        image_grid_thw = [] if image_grid_thw is None else image_grid_thw
+        video_grid_thw = [] if video_grid_thw is None else video_grid_thw
+        second_per_grid_ts = [] if second_per_grid_ts is None else second_per_grid_ts
+
+        llm_positions, mrope_position_delta = cls.get_input_positions_tensor(
+            input_tokens=input_tokens,
+            hf_config=hf_config,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            second_per_grid_ts=second_per_grid_ts,
+            context_len=context_len,
+            seq_len=seq_len,
+            audio_feature_lengths=audio_feature_lengths,
+            use_audio_in_video=use_audio_in_video,
+        )
+
+        return llm_positions.tolist(), mrope_position_delta
+
+    @classmethod
+    def get_input_positions_tensor(
+        cls,
+        input_tokens: list[int],
+        hf_config: PretrainedConfig,
+        image_grid_thw: Union[list[list[int]], torch.Tensor],
+        video_grid_thw: Union[list[list[int]], torch.Tensor],
+        second_per_grid_ts: list[float],
+        context_len: int = 0,
+        seq_len: Optional[int] = None,
+        audio_feature_lengths: Optional[torch.Tensor] = None,
+        use_audio_in_video: bool = False,
+    ) -> tuple[torch.Tensor, int]:
+        from vllm.transformers_utils.config import thinker_uses_mrope
+
+        if thinker_uses_mrope(hf_config):
+            return cls._omni_get_input_positions_tensor(
+                input_tokens=input_tokens,
+                hf_config=hf_config,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                second_per_grid_ts=second_per_grid_ts,
+                context_len=context_len,
+                seq_len=seq_len,
+                audio_feature_lengths=audio_feature_lengths,
+                use_audio_in_video=use_audio_in_video,
+            )
+        elif hf_config.model_type in ["glm4v", "glm4v_moe"]:
+            return cls._glm4v_get_input_positions_tensor(
+                input_tokens=input_tokens,
+                hf_config=hf_config,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                context_len=context_len,
+                seq_len=seq_len,
+            )
+        else:
+            return cls._vl_get_input_positions_tensor(
+                input_tokens=input_tokens,
+                hf_config=hf_config,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                second_per_grid_ts=second_per_grid_ts,
+                context_len=context_len,
+                seq_len=seq_len,
+            )
+
+    @classmethod
+    def _glm4v_get_input_positions_tensor(
+        cls,
+        input_tokens: list[int],
+        hf_config: PretrainedConfig,
+        image_grid_thw: Union[list[list[int]], torch.Tensor],
+        video_grid_thw: Union[list[list[int]], torch.Tensor],
+        context_len: int = 0,
+        seq_len: Optional[int] = None,
+    ) -> tuple[torch.Tensor, int]:
+        """Get mrope input positions and delta value for GLM4V."""
+
+        image_token_id = hf_config.image_token_id
+        video_start_token_id = hf_config.video_start_token_id
+        video_end_token_id = hf_config.video_end_token_id
+        spatial_merge_size = hf_config.vision_config.spatial_merge_size
+        llm_pos_ids_list: list = []
+
+        if not (image_grid_thw is None and video_grid_thw is None):
+            if isinstance(image_grid_thw, torch.Tensor):
+                image_grid_thw = image_grid_thw.tolist()
+
+            input_token_type: list[str] = []
+            video_check_flg = False
+            for token in input_tokens:
+                if token == video_start_token_id:
+                    video_check_flg = True
+                elif token == video_end_token_id:
+                    video_check_flg = False
+
+                if (token == image_token_id) and (video_check_flg is False):
+                    input_token_type.append("image")
+                elif (token == image_token_id) and (video_check_flg is True):
+                    input_token_type.append("video")
+                else:
+                    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]):
+                group_list = list(group_iter)
+                start_index = group_list[0][0]
+                end_index = group_list[-1][0] + 1
+                input_type_group.append((key, start_index, end_index))
+
+            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
+                if modality_type == "image":
+                    t, h, w = (
+                        image_grid_thw[mm_data_idx][0],
+                        image_grid_thw[mm_data_idx][1],
+                        image_grid_thw[mm_data_idx][2],
+                    )
+                    llm_grid_t, llm_grid_h, llm_grid_w = (
+                        t,
+                        h // spatial_merge_size,
+                        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)
+                    mm_data_idx += 1
+
+                elif modality_type == "video":
+                    t, h, w = (
+                        video_frame_num,
+                        image_grid_thw[mm_data_idx][1],
+                        image_grid_thw[mm_data_idx][2],
+                    )
+                    llm_grid_t, llm_grid_h, llm_grid_w = (
+                        t,
+                        h // spatial_merge_size,
+                        w // spatial_merge_size,
+                    )
+
+                    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)
+
+                    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)
+                    video_frame_num = 1
+
+        else:
+            text_len = len(input_tokens)
+            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1))
+
+        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+        llm_positions = llm_positions[:, context_len:seq_len]
+        mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
+        return llm_positions, mrope_position_delta
+
+    @classmethod
+    def _vl_get_input_positions_tensor(
+        cls,
+        input_tokens: list[int],
+        hf_config: PretrainedConfig,
+        image_grid_thw: Union[list[list[int]], torch.Tensor],
+        video_grid_thw: Union[list[list[int]], torch.Tensor],
+        second_per_grid_ts: list[float],
+        context_len: int = 0,
+        seq_len: Optional[int] = None,
+    ) -> tuple[torch.Tensor, int]:
+        """Get mrope input positions and delta value."""
+
+        image_token_id = hf_config.image_token_id
+        video_token_id = hf_config.video_token_id
+        vision_start_token_id = hf_config.vision_start_token_id
+        spatial_merge_size = hf_config.vision_config.spatial_merge_size
+        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_tokens = input_tokens_tensor[vision_start_indices + 1]
+        image_nums = (vision_tokens == image_token_id).sum()
+        video_nums = (vision_tokens == video_token_id).sum()
+        llm_pos_ids_list: list = []
+
+        st = 0
+        remain_images, remain_videos = image_nums, video_nums
+
+        image_index, video_index = 0, 0
+        for _ in range(image_nums + video_nums):
+            video_second_per_grid_t = 0.0
+            if image_token_id in input_tokens and remain_images > 0:
+                ed_image = input_tokens.index(image_token_id, st)
+            else:
+                ed_image = len(input_tokens) + 1
+            if video_token_id in input_tokens and remain_videos > 0:
+                ed_video = input_tokens.index(video_token_id, st)
+            else:
+                ed_video = len(input_tokens) + 1
+            if ed_image < ed_video:
+                t, h, w = (
+                    image_grid_thw[image_index][0],
+                    image_grid_thw[image_index][1],
+                    image_grid_thw[image_index][2],
+                )
+                image_index += 1
+                remain_images -= 1
+                ed = ed_image
+            else:
+                t, h, w = (
+                    video_grid_thw[video_index][0],
+                    video_grid_thw[video_index][1],
+                    video_grid_thw[video_index][2],
+                )
+                video_second_per_grid_t = 1.0
+                if second_per_grid_ts:
+                    video_second_per_grid_t = second_per_grid_ts[video_index]
+                video_index += 1
+                remain_videos -= 1
+                ed = ed_video
+
+            llm_grid_t, llm_grid_h, llm_grid_w = (
+                t,
+                h // spatial_merge_size,
+                w // spatial_merge_size,
+            )
+            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)
+
+            t_index = (
+                (
+                    torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w)
+                    * video_second_per_grid_t
+                    * tokens_per_second
+                )
+                .long()
+                .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)
+            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_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+        mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
+        llm_positions = llm_positions[:, context_len:seq_len]
+
+        return llm_positions, mrope_position_delta
+
+    @classmethod
+    def _omni_get_input_positions_tensor(
+        cls,
+        input_tokens: list[int],
+        hf_config: PretrainedConfig,
+        image_grid_thw: Union[list[list[int]], torch.Tensor],
+        video_grid_thw: Union[list[list[int]], torch.Tensor],
+        second_per_grid_ts: Optional[list[float]] = None,
+        context_len: int = 0,
+        seq_len: Optional[int] = None,
+        audio_feature_lengths: Optional[torch.Tensor] = None,
+        use_audio_in_video: bool = False,
+    ) -> tuple[torch.Tensor, int]:
+        """Get mrope input positions and delta value (Qwen2.5-Omni version).
+
+        Differences from MRotaryEmbedding:
+            1. Add audio support (and related `audio_feature_lengths`).
+            2. Add `use_audio_in_video` option to read audio from video inputs.
+                In this case, audio and vision position ids will be split into
+                chunks and interleaved.
+
+        Example:
+
+            (V_i are vision position ids, A_i are audio position ids)
+
+            |V_1 ...    V_n|A_1 ...   A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|...
+            |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
+        """
+
+        # TODO(fyabc): refactor and share more code with
+        #  _vl_get_input_positions_tensor.
+
+        thinker_config = hf_config.thinker_config
+        audio_token_id = thinker_config.audio_token_index
+        image_token_id = thinker_config.image_token_index
+        video_token_id = thinker_config.video_token_index
+        audio_start_token_id = thinker_config.audio_start_token_id
+        audio_end_token_id = thinker_config.audio_end_token_id
+        vision_start_token_id = thinker_config.vision_start_token_id
+        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)
+
+        if isinstance(image_grid_thw, list):
+            image_grid_thw = torch.tensor(image_grid_thw)
+        if isinstance(video_grid_thw, list):
+            video_grid_thw = torch.tensor(video_grid_thw)
+
+        src_item = input_tokens
+        audio_seqlens = audio_feature_lengths
+        if not second_per_grid_ts:
+            second_per_grid_ts = [1] * video_grid_thw.shape[0]
+        audio_idx = 0
+        video_idx = 0
+        image_idx = 0
+        new_src_item: list[int] = []
+        llm_pos_ids_list: list[torch.Tensor] = []
+
+        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
+            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:
+                        # 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:
+                        # processing the <|audio_bos|> after <|vision_eos|>
+                        start_idx -= 1
+                new_src_item.append(src_item[idx])
+                llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1)
+                llm_pos_ids_list.append(llm_pos_ids)
+            elif src_item[idx] == audio_token_id:
+                assert audio_seqlens is not None
+                audio_seqlen = audio_seqlens[audio_idx]
+                place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1
+                new_src_item.extend([audio_token_id] * place_num)
+                llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
+                llm_pos_ids_list.append(llm_pos_ids)
+                audio_idx += 1
+            elif src_item[idx] == image_token_id:
+                grid_t = image_grid_thw[image_idx][0]
+                grid_hs = image_grid_thw[:, 1]
+                grid_ws = image_grid_thw[:, 2]
+                t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
+                llm_pos_ids = cls._get_llm_pos_ids_for_vision(
+                    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)
+                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()
+                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)
+                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)
+                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)
+                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)
+                    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,
+                        video_idx,
+                        spatial_merge_size,
+                        t_chunk,
+                        grid_hs,
+                        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])
+                    audio_start_idx = (
+                        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)
+                            + 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)
+                    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])
+                    audio_llm_pos_ids_list = (
+                        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
+                video_idx += 1
+            # move to the next token
+            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)
+        llm_positions = llm_positions[:, context_len:seq_len]
+
+        return llm_positions, mrope_position_delta
+
+    @staticmethod
+    def _get_llm_pos_ids_for_vision(
+        start_idx: int,
+        vision_idx: int,
+        spatial_merge_size: int,
+        t_index: list[int],
+        grid_hs: torch.Tensor,
+        grid_ws: torch.Tensor,
+    ) -> torch.Tensor:
+        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()
+        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()
+        )
+        _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
+        llm_pos_ids_list.append(_llm_pos_ids + start_idx)
+        llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
+        return llm_pos_ids
+
+    @staticmethod
+    def _split_list_into_ranges(lst: torch.Tensor, interval: int) -> list[list[int]]:
+        ranges: list[list[int]] = [[] for _ in range((max(lst) // interval) + 1)]
+        for num in lst:
+            index = num // interval
+            ranges[index].append(num)
+        return ranges
+
+    @staticmethod
+    def get_next_input_positions(
+        mrope_position_delta: int,
+        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)]
+
+    @staticmethod
+    def get_next_input_positions_tensor(
+        out: np.ndarray,
+        out_offset: int,
+        mrope_position_delta: int,
+        context_len: int,
+        num_new_tokens: int,
+    ):
+        values = np.arange(
+            mrope_position_delta + context_len,
+            mrope_position_delta + context_len + num_new_tokens,
+            dtype=out.dtype,
+        )
+        out[:, out_offset : out_offset + num_new_tokens] = values
+
+    @classmethod
+    def omni_get_updates_use_audio_in_video(
+        cls,
+        thinker_config: PretrainedConfig,
+        audio_len: int,
+        video_grid_thw: Union[list[int], torch.Tensor],
+        video_second_per_grid_t: float,
+    ) -> list[int]:
+        """Get video prompt updates when `use_audio_in_video` is True.
+
+        In this case, audio and vision update ids will be split into
+        chunks and interleaved (details in `_omni_get_input_positions_tensor`).
+
+        <|video_bos|><|VIDEO|><|video_eos|> =>
+        <|video_bos|><|audio_bos|>(... chunks ...)<|audio_eos|><|video_eos|>
+        """
+
+        audio_token_id = thinker_config.audio_token_index
+        video_token_id = thinker_config.video_token_index
+        audio_start_token_id = thinker_config.audio_start_token_id
+        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)
+
+        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_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)
+            updates.extend([video_token_id] * vision_ntoken_per_chunk)
+
+            audio_chunk_size = min(t_ntoken_per_chunk, audio_len - added_audio_len)
+            updates.extend(audio_chunk_size * [audio_token_id])
+            added_audio_len += audio_chunk_size
+        if added_audio_len < audio_len:
+            updates.extend((audio_len - added_audio_len) * [audio_token_id])
+        updates.extend([audio_end_token_id])
+
+        return updates

vllm_omni/model_executor/model_loader/weight_utils.py

@@ -0,0 +1,67 @@
+import time
+from pathlib import Path
+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
+
+logger = init_logger(__name__)
+
+
+def download_weights_from_hf_specific(
+    model_name_or_path: str,
+    cache_dir: Optional[str],
+    allow_patterns: list[str],
+    revision: Optional[str] = None,
+    ignore_patterns: Optional[Union[str, list[str]]] = None,
+) -> str:
+    """Download model weights from Hugging Face Hub. Users can specify the
+    allow_patterns to download only the necessary weights.
+
+    Args:
+        model_name_or_path (str): The model name or path.
+        cache_dir (Optional[str]): The cache directory to store the model
+            weights. If None, will use HF defaults.
+        allow_patterns (list[str]): The allowed patterns for the
+            weight files. Files matched by any of the patterns will be
+            downloaded.
+        revision (Optional[str]): The revision of the model.
+        ignore_patterns (Optional[Union[str, list[str]]]): The patterns to
+            filter out the weight files. Files matched by any of the patterns
+            will be ignored.
+
+    Returns:
+        str: The path to the downloaded model weights.
+    """
+    assert len(allow_patterns) > 0
+    local_only = huggingface_hub.constants.HF_HUB_OFFLINE
+
+    logger.info("Using model weights format %s", allow_patterns)
+    # Use file lock to prevent multiple processes from
+    # downloading the same model weights at the same time.
+    with get_lock(model_name_or_path, cache_dir):
+        start_time = time.perf_counter()
+        for allow_pattern in allow_patterns:
+            hf_folder = snapshot_download(
+                model_name_or_path,
+                allow_patterns=allow_pattern,
+                ignore_patterns=ignore_patterns,
+                cache_dir=cache_dir,
+                tqdm_class=DisabledTqdm,
+                revision=revision,
+                local_files_only=local_only,
+            )
+            # If we have downloaded weights for this allow_pattern,
+            # we don't need to check the rest.
+            if any(Path(hf_folder).glob(allow_pattern)):
+                break
+        time_taken = time.perf_counter() - start_time
+        if time_taken > 0.5:
+            logger.info(
+                "Time spent downloading weights for %s: %.6f seconds",
+                model_name_or_path,
+                time_taken,
+            )
+    return hf_folder

vllm_omni/model_executor/models/__init__.py

@@ -0,0 +1 @@
+from .registry import OmniModelRegistry  # noqa: F401

vllm_omni/model_executor/models/qwen2_5_omni.py

@@ -0,0 +1,871 @@
+import glob
+import os
+from collections.abc import Iterable
+from functools import cached_property
+from typing import NamedTuple, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
+    Qwen2_5OmniConfig,
+    Qwen2_5OmniTalkerConfig,
+    Qwen2_5OmniThinkerConfig,
+)
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerDummyInputsBuilder,
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    Qwen2_5OmniThinkerProcessingInfo,
+)
+from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
+
+# from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
+from vllm.multimodal import MULTIMODAL_REGISTRY
+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.models.utils import add_prefix_to_loaded_weights
+
+TALKER_CODEC_EOS_TOKEN_ID = 8294
+TALKER_CODEC_BOS_TOKEN_ID = 8293
+
+
+class OmniOutput(NamedTuple):
+    """Output from the merged Omni model containing both text and audio."""
+
+    text_hidden_states: torch.Tensor
+    multimodal_outputs: dict = {}
+    intermediate_tensors: Optional[IntermediateTensors] = None
+
+
+logger = init_logger(__name__)
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    info=Qwen2_5OmniThinkerProcessingInfo,
+    dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
+)
+class Qwen2_5OmniForConditionalGeneration(
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.have_multimodal_outputs = True
+        config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+        # keep vllm_config for later submodule init
+        self.vllm_config = vllm_config
+
+        # Initialize thinker components
+        thinker_config: Qwen2_5OmniThinkerConfig = config.thinker_config
+        self.thinker_config = thinker_config
+        self.multimodal_config = multimodal_config
+
+        # Initialize talker components
+        talker_config: Qwen2_5OmniTalkerConfig = config.talker_config
+        self.talker_config = talker_config
+
+        self.model_stage = vllm_config.model_config.model_stage
+        if self.model_stage == "thinker":
+            # Initialize thinker model (multimodal processing)
+            self.thinker = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "thinker"),
+                hf_config=thinker_config,
+                # Use registry architecture key
+                architectures=["Qwen2_5OmniThinkerModel"],
+            )
+            self.model = self.thinker
+            self.talker = None
+            self.token2wav = None
+
+        elif self.model_stage == "talker":
+            self.thinker = None
+            # Initialize talker model wrapper (handles projection + LM)
+            self.talker = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "talker"),
+                hf_config=talker_config,
+                # Use registry architecture key
+                architectures=["Qwen2_5OmniTalkerModel"],
+            )
+            self.talker.init_multi_modal(thinker_config)
+            self.model = self.talker
+            self.token2wav = None
+
+        elif self.model_stage == "code2wav":
+            self.thinker = None
+            self.talker = None
+            # Initialize token2wav (code->mel->wav) like thinker/talker
+            self.token2wav_config = getattr(config, "token2wav_config", None)
+            self.token2wav = None
+            if self.token2wav_config is not None:
+                self.token2wav = init_vllm_registered_model(
+                    vllm_config=vllm_config,
+                    prefix=maybe_prefix(prefix, "token2wav"),
+                    hf_config=self.token2wav_config,
+                    architectures=["Qwen2_5OmniToken2WavModel"],
+                )
+            # voice resources (loaded on demand)
+            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
+        )
+
+    # -------------------- Device utilities --------------------
+    @staticmethod
+    def _module_device(module: nn.Module) -> torch.device:
+        try:
+            return next(module.parameters()).device
+        except StopIteration:
+            # No parameters; fall back to buffers or cpu
+            for _, buf in module.named_buffers(recurse=True):
+                return buf.device
+            return torch.device("cpu")
+
+    def move_submodules_to_devices(
+        self,
+        *,
+        thinker_device: Optional[Union[str, torch.device]] = None,
+        talker_device: Optional[Union[str, torch.device]] = None,
+        token2wav_device: Optional[Union[str, torch.device]] = None,
+    ) -> None:
+        """Optionally move thinker/talker/token2wav to different devices.
+
+        Example:
+            model.move_submodules_to_devices(
+                thinker_device='cuda:0',
+                talker_device='cuda:1',
+                token2wav_device='cpu',
+            )
+        """
+        if thinker_device is not None and self.thinker is not None:
+            self.thinker.to(thinker_device)
+        if talker_device is not None and self.talker is not None:
+            self.talker.to(talker_device)
+        if token2wav_device is not None and self.token2wav is not None:
+            self.token2wav.to(token2wav_device)
+
+    @cached_property
+    def sampler(self):
+        if hasattr(self.model, "sampler"):
+            return self.model.sampler
+        return Sampler()
+
+    def get_input_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        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 self.model.get_input_embeddings(input_ids, multimodal_embeddings)
+
+    def get_multimodal_embeddings(self, **kwargs):
+        # Delegate to thinker model for multimodal processing
+        return self.model.get_multimodal_embeddings(**kwargs)
+
+    def last_index_of(self, list, value):
+        return len(list) - 1 - list[::-1].index(value)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        generate_audio: bool = True,
+        voice_type: str = "Chelsie",
+        codec: Optional[torch.Tensor] = None,
+        sampling_metadata: Optional[SamplingMetadata] = None,
+        logits_index: Optional[int] = None,
+        sampler=None,
+        additional_information: Optional[dict[str, object]] = None,
+        **kwargs: object,
+    ) -> Union[torch.Tensor, IntermediateTensors, OmniOutput]:
+        """
+        Workflow:
+        1) Thinker: multimodal understanding → text hidden states.
+        2) If audio requested and codec not provided, use talker to derive codec.
+        3) If audio requested (or codec provided), use token2wav to synthesize waveform.
+        4) Return text hidden states (and audio when applicable).
+        """
+        if self.model_stage == "thinker":
+            # Normalize to batched inputs if caller provides 1D/2D unbatched tensors
+            added_batch_dim = False
+            if input_ids is not None and input_ids.ndim == 1:
+                input_ids = input_ids.unsqueeze(0)
+                added_batch_dim = True
+            if positions is not None and positions.ndim == 1:
+                positions = positions.unsqueeze(0)
+                added_batch_dim = True
+            if inputs_embeds is not None and inputs_embeds.ndim == 2:
+                inputs_embeds = inputs_embeds.unsqueeze(0)
+                added_batch_dim = True
+            thinker_dev = self._module_device(self.thinker)
+
+            # 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(
+                    0
+                )  # (1, 0)
+                added_batch_dim = True
+
+            # 1) Thinker (ensure inputs on thinker's device)
+            if input_ids is not None and input_ids.device != thinker_dev:
+                input_ids = input_ids.to(thinker_dev)
+            if positions is not None and positions.device != thinker_dev:
+                positions = positions.to(thinker_dev)
+            if inputs_embeds is not None and inputs_embeds.device != thinker_dev:
+                inputs_embeds = inputs_embeds.to(thinker_dev)
+            # Run thinker
+            thinker_output = self.thinker(
+                input_ids=input_ids,
+                positions=positions[0],
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+                **kwargs,
+            )
+
+            if isinstance(thinker_output, tuple):
+                embeds, text_hidden_states = thinker_output
+            else:
+                text_hidden_states = thinker_output
+
+            # Text-only path
+            return OmniOutput(
+                text_hidden_states=(text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states),
+                multimodal_outputs=None,
+            )
+
+        # 2) Talker (if codec not provided)
+        if self.model_stage == "talker":
+            # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
+            # 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),
+            #   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.
+
+            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,
+                )
+                additional_information = {}
+                self.thinker_reply_part = torch.zeros_like(inputs_embeds)
+                is_profile = True
+            else:
+                is_profile = False
+
+            # Ensure we have base embeddings when only ids are provided
+            if inputs_embeds is None and input_ids is not None:
+                inputs_embeds = self.talker.get_input_embeddings(input_ids)
+
+            # ------- 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")
+            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)
+            ):
+                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]
+            elif isinstance(request_ids, list) and isinstance(runtime_addi, dict):
+                for rid in request_ids:
+                    if isinstance(rid, str) and isinstance(runtime_addi.get(rid), dict):
+                        runtime_addi_by_req[rid] = runtime_addi[rid]
+
+            # Containers to return per-request updates (e.g., thinker_reply_part_per_request)
+            update_by_req_id: dict[str, dict] = {}
+
+            # ------- Prefill: span_len > 1 -------
+            if (
+                not is_profile
+                and isinstance(request_ids, list)
+                and isinstance(request_token_spans, list)
+                and isinstance(addi_by_req, dict)
+            ):
+                for idx_req, rid in enumerate(request_ids):
+                    s, e = request_token_spans[idx_req]
+                    span_len = int(e) - int(s)
+                    if span_len <= 1:
+                        continue
+                    info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {}
+                    if not isinstance(info, dict):
+                        info = {}
+                    pe = info.get("prompt_embeds")  # Tensor [P,H]
+                    tr = info.get("thinker_result")  # Tensor [K,H]
+                    ptoks = info.get("prompt_token_ids")  # list[int]
+                    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),
+                        )
+                    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),
+                        )
+                    if not isinstance(ptoks, (list, torch.Tensor)):
+                        ptoks = []
+                    if not isinstance(otoks, (list, torch.Tensor)):
+                        otoks = []
+
+                    req_input_ids, req_embeds = self._thinker_to_talker_prefill(
+                        voice_type=voice_type,
+                        output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                        output_token_ids=otoks,
+                        thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                        prompt_token_ids=ptoks,
+                    )
+                    seg_len = min(span_len, req_embeds.shape[0])
+                    inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
+                    if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
+                        input_ids[s : s + seg_len] = req_input_ids
+
+                    # 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()
+                        )
+
+            # ------- Decode: span_len == 1 -------
+            if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list):
+                for idx_req, rid in enumerate(request_ids):
+                    s, e = request_token_spans[idx_req]
+                    if (int(e) - int(s)) != 1:
+                        continue
+                    # choose step vector in priority order
+                    step_vec = None
+                    # A) runtime queue
+                    q = None
+                    if isinstance(rid, str):
+                        q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request")
+                    if isinstance(q, torch.Tensor) and q.numel() > 0:
+                        step_vec = q[0:1]
+                        new_q = q[1:].detach().to("cpu").contiguous()
+                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q
+                    else:
+                        # B) per-request provided decode vector (optional)
+                        info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {}
+                        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
+                        ):
+                            # C) fallback shared pool
+                            step_vec = self.thinker_reply_part[0:1]
+                            self.thinker_reply_part = self.thinker_reply_part[1:]
+
+                    if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
+                        one_id = input_ids[s : s + 1]
+                        _, one_embed = self._thinker_to_talker_decode_one_step(
+                            output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                            output_token_ids=one_id,
+                        )
+                        inputs_embeds[s] = one_embed[0]
+
+            with torch.inference_mode():
+                talker_hidden = self.talker(
+                    input_ids=input_ids,
+                    positions=positions[0],
+                    inputs_embeds=inputs_embeds,
+                )
+            multimodal_outputs: dict = None
+            # Return updates if any
+            if update_by_req_id:
+                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
+                sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448
+
+            return OmniOutput(
+                text_hidden_states=talker_hidden,
+                multimodal_outputs=multimodal_outputs,
+            )
+
+        if self.model_stage == "code2wav":
+            code = (
+                input_ids
+                if input_ids is not None
+                else torch.zeros(
+                    inputs_embeds.shape[0],
+                    dtype=torch.long,
+                    device=inputs_embeds.device,
+                )
+            )
+
+            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})
+
+        return OmniOutput(
+            text_hidden_states=torch.cat(
+                [
+                    torch.zeros(
+                        [inputs_embeds.shape[0], self.talker.config.hidden_size],
+                        dtype=torch.bfloat16,
+                    ).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])
+                            .to(torch.bfloat16)
+                            .to(self._module_device(self.model))
+                        )
+                    )[0],
+                ],
+                dim=0,
+            ),
+            multimodal_outputs=None,
+        )
+
+    def generate_audio(self, code, voice_type):
+        token2wav_dev = self._module_device(self.token2wav)
+        if isinstance(code, torch.Tensor):
+            code_tensor = code.to(dtype=torch.long, device=token2wav_dev)
+        else:
+            code_tensor = torch.as_tensor(code, dtype=torch.long, device=token2wav_dev)
+        if code_tensor.ndim == 2 and code_tensor.shape[0] == 1:
+            code_tensor = code_tensor.squeeze(0)
+
+        audio_tensor = self._codec_to_audio(code_tensor, voice_type)
+
+        return audio_tensor
+
+    def _load_talker_embedding(
+        self,
+    ) -> torch.nn.Embedding:
+        return self.talker.language_model.model.embed_tokens
+
+    def _init_special_tokens_embeddings(
+        self,
+    ):
+        # talker embeddings
+        self.talker_embedding = self._load_talker_embedding()
+
+        # embed_text_bos_token
+        self.tts_text_spk_token_ids = {
+            # M02: Male voice with standard Mandarin and a slight northern accent
+            "m02": 151870,
+            "Ethan": 151870,
+            # F030: Your anime-styled virtual girlfriend
+            "f030": 151872,
+            "Chelsie": 151872,
+        }
+        self.default_tts_text_spk_type = list(self.tts_text_spk_token_ids.keys())[0]
+        self.tts_text_spk_token_ids["prefix_caching"] = 151870
+
+        talker_hf_config = self.talker_config
+        if hasattr(talker_hf_config, "talker_config"):
+            talker_hf_config = talker_hf_config.talker_config
+
+        self.embed_text_bos_token = self.thinker_embedding(
+            torch.tensor(
+                [talker_hf_config.tts_text_start_token_id],
+                dtype=torch.long,
+                device=self._module_device(self.talker),
+            )
+        )
+        self.embed_text_spk_tokens = {
+            key: self.thinker_embedding(
+                torch.tensor(
+                    [value],
+                    dtype=torch.long,
+                    device=self._module_device(self.talker),
+                )
+            )
+            for key, value in self.tts_text_spk_token_ids.items()
+        }
+        self.embed_text_eos_token = self.thinker_embedding(
+            torch.tensor(
+                [talker_hf_config.tts_text_end_token_id],
+                dtype=torch.long,
+                device=self._module_device(self.talker),
+            )
+        )
+        self.embed_text_pad_token = self.thinker_embedding(
+            torch.tensor(
+                [talker_hf_config.tts_text_pad_token_id],
+                dtype=torch.long,
+                device=self._module_device(self.talker),
+            )
+        )
+        self.embed_codec_bos_token = self.talker_embedding(
+            torch.tensor(
+                [talker_hf_config.tts_codec_start_token_id],
+                dtype=torch.long,
+                device=self._module_device(self.talker),
+            )
+        )
+        self.embed_codec_eos_token = self.talker_embedding(
+            torch.tensor(
+                [talker_hf_config.tts_codec_end_token_id],
+                dtype=torch.long,
+                device=self._module_device(self.talker),
+            )
+        )
+        self.embed_codec_pad_token = self.talker_embedding(
+            torch.tensor(
+                [talker_hf_config.tts_codec_pad_token_id],
+                dtype=torch.long,
+                device=self._module_device(self.talker),
+            )
+        )
+        return set(["thinker_embedding.weight", "talker_embedding.weight"])
+
+    def _get_embed_text_spk_token(self, voice_type: str):
+        if voice_type not in self.embed_text_spk_tokens:
+            return self.embed_text_bos_token
+        return self.embed_text_spk_tokens[voice_type]
+
+    def _get_text_spk_token_id(self, voice_type: str):
+        talker_hf_config = self.talker_config
+        if hasattr(talker_hf_config, "talker_config"):
+            talker_hf_config = talker_hf_config.talker_config
+
+        if voice_type not in self.tts_text_spk_token_ids:
+            return talker_hf_config.tts_text_start_token_id
+        return self.tts_text_spk_token_ids[voice_type]
+
+    def _thinker_to_talker_prefill(
+        self,
+        voice_type: str,
+        output_prompt_embeds,
+        output_token_ids,
+        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
+
+        # if len(output.outputs[0].token_ids) == 2:
+        # issue request
+        prompt_embeds = torch.cat(
+            [
+                thinker_prompt_embeds,
+                self._get_embed_text_spk_token(voice_type) + self.embed_codec_pad_token,
+                output_prompt_embeds[:1] + self.embed_codec_bos_token,
+            ],
+            dim=0,
+        )
+
+        prompt_token_ids_processed = prompt_token_ids + [
+            talker_hf_config.tts_codec_pad_token_id,
+            output_token_ids[0],
+        ]
+        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) + [
+                self.talker_config.tts_codec_pad_token_id,
+                self.talker_config.tts_codec_start_token_id,
+            ]
+        else:
+            prompt_token_ids_processed = [
+                self.talker_config.tts_codec_pad_token_id,
+                self.talker_config.tts_codec_start_token_id,
+            ][-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))
+            )
+        return prompt_token_ids_processed, prompt_embeds
+
+    def _thinker_to_talker_decode_one_step(
+        self,
+        output_prompt_embeds,
+        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]:
+        # Handle OmniOutput type
+        if isinstance(hidden_states, OmniOutput):
+            hidden_states = hidden_states.text_hidden_states
+
+        # Use thinker model for logits computation
+        return self.model.compute_logits(hidden_states)
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[SamplerOutput]:
+        # Use thinker model for sampling
+        return self.model.sample(logits, sampling_metadata)
+
+    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.
+
+        Args:
+            text_tokens: Text tokens from thinker model
+            voice_type: Voice type for speech generation
+
+        Returns:
+            Audio tensor
+        """
+        # Generate codec tokens using talker model
+        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)
+
+        # Generate audio using token2wav model
+        return self._codec_to_audio(codec_tokens, voice_type=voice_type)
+
+    def _convert_to_codec_tokens(
+        self, talker_output: torch.Tensor, sampling_metadata: SamplingMetadata
+    ) -> torch.Tensor:
+        """
+        Reference (HF): use the talker's codec head to obtain logits, suppress BOS,
+        then greedily select the next codec token for the current step.
+        """
+        with torch.inference_mode():
+            logits = self.talker.compute_logits(talker_output, None)
+            if logits is None:
+                return torch.zeros(
+                    (talker_output.size(0), 0),
+                    dtype=torch.long,
+                    device=talker_output.device,
+                )
+
+            # 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"):
+                bos_id = int(getattr(self.talker_config, "tts_codec_start_token_id"))
+            if bos_id is not None:
+                logits[..., bos_id] = -1e9
+
+            # Take the distribution at the last step and select greedily
+            next_id = self.talker.sample(logits, sampling_metadata).sampled_token_ids
+            return next_id.to(dtype=torch.long)
+
+    def _init_token2wav_model(self, hf_model_folder):
+        """Initialize speaker resources if provided; model is constructed in
+        __init__."""
+        if self.token2wav is None or self.token2wav_config is None:
+            return
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        # optional speaker resources
+        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()}
+        # legacy: load from directory if provided
+        model_path = hf_model_folder
+        if isinstance(model_path, str) and os.path.isdir(model_path):
+            spk_pt = os.path.join(model_path, "spk_dict.pt")
+            if os.path.exists(spk_pt):
+                data = torch.load(spk_pt, map_location=device)
+                for key, value in data.items():
+                    self._token2wav_conds[key] = value["cond"].to(device)
+                    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"))):
+                    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"))):
+                    key = os.path.basename(f).split("_")[0].lower()
+                    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]:
+        if self.token2wav is None:
+            self._init_token2wav_model()
+        if self.token2wav is None:
+            return None
+        # Normalize voice type
+        voice = voice_type or "default"
+        # Resolve cond / ref_mel if provided
+        cond = None
+        ref_mel = None
+        if voice in self._token2wav_conds and voice in self._token2wav_ref_mels:
+            cond = self._token2wav_conds[voice]
+            ref_mel = self._token2wav_ref_mels[voice]
+        # Fallback: create dummy cond/ref_mel if not provided
+        token2wav_dev = self._module_device(self.token2wav)
+        if cond is None:
+            cond = torch.zeros(
+                (1, self.token2wav_config.dit_config.enc_emb_dim),
+                device=token2wav_dev,
+                dtype=torch.float32,
+            )
+        if ref_mel is None:
+            ref_mel = torch.zeros(
+                (1, 300, self.token2wav_config.dit_config.mel_dim),
+                device=token2wav_dev,
+                dtype=torch.float32,
+            )
+
+        # Ensure codec is (1, T) long tensor on correct device
+        if isinstance(codec_tokens, torch.Tensor):
+            codec = codec_tokens.to(dtype=torch.long, device=token2wav_dev)
+            if codec.ndim == 1:
+                codec = codec.unsqueeze(0)
+        else:
+            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)
+        factor = getattr(self.token2wav.token2wav.factor, "factor", 2)
+        chunk_size = 48
+        mel_dim = getattr(
+            self.token2wav.token2wav.code2wav_dit_model,
+            "mel_dim",
+            self.token2wav_config.dit_config.mel_dim,
+        )
+        total_mel = int(codec.shape[1] * factor)
+        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)
+
+        logger.info(
+            "Currently, we do not use the chunked process, we only use the "
+            "token2wav.process_chunk for the whole sequence. "
+            "The stream mode will be implemented in the future."
+        )
+
+        chunk_ends = []
+        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:
+                chunk_ends.append(i)
+
+        # Number of chunks in mel domain
+        prev_generated = None
+        wav_chunks: list = []
+
+        with torch.inference_mode():
+            for n, i in enumerate([0]):
+                finished = i == codec.shape[1] - 1
+                _, audio_chunk = self.token2wav.process_chunk(
+                    conditioning=cond,
+                    reference_mel=ref_mel,
+                    codec_all=codec,
+                    y_all=y_all,
+                    i=n,
+                    steps=steps,
+                    prev_generated=prev_generated if prev_generated is not None else [],
+                    finished=True,
+                )
+                prev_generated = audio_chunk
+                wav_chunks.append(audio_chunk.detach().cpu().numpy())
+
+        if len(wav_chunks) == 0:
+            return torch.zeros(0, device=token2wav_dev)
+
+        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]:
+        """Load weights for all components of the omni model."""
+        loaded_weights = set()
+        thinker_weights = []
+        talker_weights = []
+        token2wav_weights = []
+        for k, v in weights:
+            if k.startswith("thinker."):
+                thinker_weights.append((k, v))
+            elif k.startswith("talker."):
+                talker_weights.append((k, v))
+            elif k.startswith("token2wav."):
+                token2wav_weights.append((k, v))
+            else:
+                raise ValueError(f"Unknown weight prefix: {k}")
+
+        # Load thinker weights
+        if self.thinker:
+            if thinker_weights:
+                thinker_loaded = self.thinker.load_weights(thinker_weights)
+            else:
+                thinker_loaded = set([k for k, v in thinker_weights])
+            thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker")
+            loaded_weights.update(thinker_loaded)
+
+        # Load talker weights
+        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"]
+                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))
+                    )
+            talker_loaded = self.talker.load_weights(talker_weights)
+            talker_loaded = add_prefix_to_loaded_weights(talker_loaded, "talker")
+            loaded_weights.update(talker_loaded)
+            loaded_weights.update(self._init_special_tokens_embeddings())
+
+        # Load token2wav weights (if any)
+        if token2wav_weights and self.token2wav is not None:
+            # download weights from huggingface for spk_dict.pt
+            model_path = self.vllm_config.model_config.model
+            download_dir = self.vllm_config.load_config.download_dir
+            if os.path.exists(model_path):
+                hf_model_folder = model_path
+            else:
+                hf_model_folder = download_weights_from_hf_specific(
+                    model_path,
+                    download_dir,
+                    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 = add_prefix_to_loaded_weights(t2w_loaded, "token2wav")
+            loaded_weights.update(t2w_loaded)
+
+        return loaded_weights

vllm_omni/model_executor/models/qwen2_5_omni_talker.py

@@ -0,0 +1,239 @@
+from collections.abc import Iterable
+from functools import cached_property
+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 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.qwen2_5_omni_thinker import (
+    Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerDummyInputsBuilder,
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    Qwen2_5OmniThinkerProcessingInfo,
+)
+from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    init_vllm_registered_model,
+    maybe_prefix,
+    merge_multimodal_embeddings,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+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
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    info=Qwen2_5OmniThinkerProcessingInfo,
+    dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
+)
+class Qwen2_5OmniTalkerForConditionalGeneration(
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+):
+    logger = init_logger(__name__)
+    # Align to thinker-style static mapper for clarity
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            # text LM head/body in talker
+            "talker.codec_head.": "language_model.lm_head.",
+            "talker.model.": "language_model.model.",
+            # projection weights
+            "talker.thinker_to_talker_proj.": "thinker_to_talker_proj.",
+            # fallback root
+            "talker.": "",
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config: Qwen2_5OmniTalkerConfig = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+
+        self.prefix = prefix
+        self.quant_config = quant_config
+
+        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
+        else:
+            self.config = config
+
+        self.thinker_to_talker_proj = ColumnParallelLinear(
+            self.config.embedding_size,
+            self.config.hidden_size,
+            bias=True,
+            gather_output=True,
+            skip_bias_add=False,
+            quant_config=quant_config,
+        )
+        self.language_model = init_vllm_registered_model(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "language_model"),
+            hf_config=getattr(self.config, "text_config", self.config),
+            architectures=["Qwen2ForCausalLM_old"],
+        )
+        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)
+        self.visual = Qwen2_5_VisionTransformer(
+            vision_config=thinker_config.vision_config,
+            norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
+            quant_config=self.quant_config,
+            prefix=maybe_prefix(self.prefix, "visual"),
+        )
+
+    def get_language_model(self) -> torch.nn.Module:
+        return self.language_model
+
+    @cached_property
+    def sampler(self):
+        if hasattr(self.language_model, "sampler"):
+            return self.language_model.sampler
+
+        return Sampler()
+
+    def get_input_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
+    ) -> 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(
+                input_ids,
+                inputs_embeds,
+                multimodal_embeddings,
+                [
+                    self.config.image_token_index,
+                    self.config.video_token_index,
+                    self.config.audio_token_index,
+                ],
+            )
+        return inputs_embeds
+
+    def forward(
+        self,
+        input_ids: torch.Tensor = None,
+        positions: torch.Tensor = None,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        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"
+        # forward_context: ForwardContext = get_forward_context()  # unused variable
+
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+        elif inputs_embeds is None:
+            # for profile_run:
+            inputs_embeds = self.get_input_embeddings(input_ids)
+
+        input_ids = None
+
+        # projection
+        inputs_embeds, _ = self.thinker_to_talker_proj(inputs_embeds)
+
+        hidden_states = self.language_model.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
+        )
+        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]:
+        logits = self.language_model.compute_logits(hidden_states)
+        logits = self.bad_word_processor(logits)
+        return logits
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[SamplerOutput]:
+        return self.language_model.sample(logits, sampling_metadata)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=["thinker.", "token2wav."],
+        )
+        loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        # Log load summary
+        try:
+            total_bytes = 0
+            for name, param in self.named_parameters():
+                if param is not None and param.data is not None:
+                    total_bytes += param.data.numel() * param.data.element_size()
+            device = next(self.parameters()).device
+            self.logger.info(
+                "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
+                self.__class__.__name__,
+                True,
+                total_bytes / (1024**2),
+                str(device),
+            )
+        except Exception:
+            pass
+        multi_model_weights = set()
+        for name, param in self.visual.named_parameters():
+            multi_model_weights.add("visual." + name)
+        for name, param in self.audio_tower.named_parameters():
+            multi_model_weights.add("audio_tower." + name)
+        loaded.update(multi_model_weights)
+        return loaded
+
+    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
+        mm_input_by_modality = {}
+
+        # 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)
+        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 []
+
+        # The result multimodal_embeddings is tuple of tensors, with each
+        # tensor correspoending to a multimodal data item (image or video).
+        multimodal_embeddings: tuple[torch.Tensor, ...] = ()
+
+        # NOTE: It is important to iterate over the keys in this dictionary
+        # to preserve the order of the modalities.
+        for modality in mm_input_by_modality:
+            multimodal_input = mm_input_by_modality[modality]
+            if modality == "image":
+                vision_embeddings = self._process_image_input(multimodal_input)
+                multimodal_embeddings += vision_embeddings
+            if modality == "video":
+                video_embeddings = self._process_video_input(multimodal_input)
+                multimodal_embeddings += video_embeddings
+            if modality == "audio":
+                audio_embeddings = self._process_audio_input(multimodal_input)
+                multimodal_embeddings += audio_embeddings
+        return multimodal_embeddings

vllm_omni/model_executor/models/qwen2_5_omni_thinker.py

@@ -0,0 +1,840 @@
+from collections.abc import Iterable, Mapping, Sequence
+from copy import copy
+from functools import partial
+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.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.qwen2_5_vl import (
+    Qwen2_5_VisionTransformer,
+    Qwen2_5_VLImageEmbeddingInputs,
+    Qwen2_5_VLImageInputs,
+    Qwen2_5_VLImagePixelInputs,
+    Qwen2_5_VLProcessingInfo,
+    Qwen2_5_VLVideoEmbeddingInputs,
+    Qwen2_5_VLVideoInputs,
+    Qwen2_5_VLVideoPixelInputs,
+)
+from vllm.model_executor.models.qwen2_audio import (
+    Qwen2AudioInputs,
+    Qwen2AudioProcessingInfo,
+    _get_feat_extract_output_lengths,
+)
+from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    init_vllm_registered_model,
+    maybe_prefix,
+    merge_multimodal_embeddings,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import (
+    ImageItem,
+    ModalityData,
+    MultiModalDataDict,
+    MultiModalFieldConfig,
+    MultiModalKwargs,
+    NestedTensors,
+)
+from vllm.multimodal.parse import (
+    AudioProcessorItems,
+    DictEmbeddingItems,
+    ModalityDataItems,
+    MultiModalDataItems,
+    MultiModalDataParser,
+)
+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:
+    import flash_attn
+except (ImportError, ModuleNotFoundError):
+    flash_attn = None
+
+logger = init_logger(__name__)
+
+
+def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
+    audio_feature_lengths = hf_inputs.get("audio_feature_lengths", torch.empty((0,)))
+
+    image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
+    image_grid_sizes = image_grid_thw.prod(-1)
+
+    video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3)))
+    video_grid_sizes = video_grid_thw.prod(-1)
+
+    return dict(
+        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),
+        video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes),
+        video_grid_thw=MultiModalFieldConfig.batched("video"),
+        second_per_grid_ts=MultiModalFieldConfig.batched("video"),
+    )
+
+
+class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
+    def _parse_audio_data(
+        self,
+        data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],
+    ) -> ModalityDataItems[Any, Any]:
+        if isinstance(data, dict):
+            return DictEmbeddingItems(
+                data,
+                modality="audio",
+                required_fields={"input_audio_features", "audio_feature_lengths"},
+                fields_factory=_qwen2_5_omni_thinker_field_config,
+            )
+
+        return super()._parse_audio_data(data)
+
+
+class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.get_hf_config(Qwen2_5OmniConfig).thinker_config
+
+    def get_hf_processor(self, **kwargs: object) -> Qwen2_5OmniProcessor:
+        return self.ctx.get_hf_processor(
+            Qwen2_5OmniProcessor,
+            use_fast=kwargs.pop("use_fast", True),
+            **kwargs,
+        )
+
+    def get_feature_extractor(self, **kwargs: object):
+        hf_processor = self.get_hf_processor(**kwargs)
+        feature_extractor = hf_processor.feature_extractor  # type: ignore
+        assert isinstance(feature_extractor, WhisperFeatureExtractor)
+        return feature_extractor
+
+    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
+        return {"audio": None, "image": None, "video": None}
+
+
+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)
+        num_videos = mm_counts.get("video", 0)
+
+        hf_processor = self.info.get_hf_processor()
+
+        audio_token: str = hf_processor.audio_token
+        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
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> MultiModalDataDict:
+        num_audios = mm_counts.get("audio", 0)
+        num_images = mm_counts.get("image", 0)
+        num_videos = mm_counts.get("video", 0)
+
+        feature_extractor = self.info.get_feature_extractor()
+
+        target_audio_length = (
+            min(
+                feature_extractor.chunk_length,
+                30,
+            )
+            * 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)
+
+        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),
+            "video": self._get_dummy_videos(
+                width=target_width,
+                height=target_height,
+                num_frames=target_num_frames,
+                num_videos=num_videos,
+            ),
+        }
+
+        return mm_data
+
+
+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)
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        mm_data = dict(mm_data)
+        audios = mm_data.pop("audios", [])
+
+        # NOTE: WhisperFeatureExtractor cannot handle empty list of audios
+        if audios:
+            # NOTE: Qwen2.5-Omni processor accept "audio"
+            mm_data["audio"] = audios
+            mm_kwargs = dict(
+                **mm_kwargs,
+            )
+
+        hf_inputs = super()._call_hf_processor(
+            prompt=prompt,
+            mm_data=mm_data,
+            mm_kwargs=mm_kwargs,
+            tok_kwargs=tok_kwargs,
+        )
+
+        input_features = hf_inputs.pop("input_features", None)
+        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)
+            hf_inputs["input_audio_features"] = input_features
+        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
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return _qwen2_5_omni_thinker_field_config(hf_inputs)
+
+    def _maybe_apply_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        prompt_ids: list[int],
+        mm_kwargs: MultiModalKwargs,
+        is_update_applied: bool,
+    ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]:
+        """
+        Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
+        """
+        unbound_prompt_updates = self._get_prompt_updates(
+            mm_items,
+            hf_processor_mm_kwargs,
+            mm_kwargs,
+        )
+        mm_prompt_updates = self._bind_and_group_updates(unbound_prompt_updates)
+
+        mm_item_counts = mm_items.get_all_counts()
+        self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
+
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
+
+        if is_update_applied:
+            mm_placeholders = self._find_mm_placeholders(
+                mm_prompt_updates,
+                prompt_ids,
+                mm_item_counts,
+            )
+            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)
+        else:
+            (
+                prompt_ids,
+                prompt,
+                mm_placeholders,
+            ) = self._apply_prompt_updates(
+                prompt_ids,
+                mm_prompt_updates,
+                mm_item_counts,
+            )
+            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)
+
+        if use_audio_in_video:
+            mm_kwargs["use_audio_in_video"] = True
+
+        return prompt_ids, prompt, mm_placeholders
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, Any],
+        out_mm_kwargs: MultiModalKwargs,
+    ) -> Sequence[PromptUpdate]:
+        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        tokenizer = self.info.get_tokenizer()
+        image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)
+        vocab = tokenizer.get_vocab()
+
+        audio_token = processor.audio_token
+        image_token = processor.image_token
+        video_token = processor.video_token
+        audio_token_id = vocab[audio_token]
+        image_token_id = vocab[image_token]
+        video_token_id = vocab[video_token]
+
+        audio_feature_lengths = out_mm_kwargs.get("audio_feature_lengths")
+        feature_attention_mask = out_mm_kwargs.get("feature_attention_mask")
+        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_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_lengths = audio_output_lens.tolist()
+
+        # number of audios read from video.
+        audio_in_video_item_idx = 0
+
+        def get_replacement_qwen2_audio(item_idx: int):
+            item_idx += audio_in_video_item_idx
+
+            num_features = audio_output_lengths[item_idx]
+            if num_features == 0:
+                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"
+                )
+
+            return [audio_token_id] * num_features
+
+        def get_replacement_qwen2_vision(item_idx: int, modality: str):
+            grid_thw = out_mm_kwargs[f"{modality}_grid_thw"][item_idx]
+            assert isinstance(grid_thw, torch.Tensor)
+            merge_length = image_processor.merge_size**2
+
+            token_id = image_token_id if modality == "image" else video_token_id
+            return [token_id] * (int(grid_thw.prod()) // merge_length)
+
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
+        thinker_config = self.info.get_hf_config()
+
+        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]
+            video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx]
+
+            audio_in_video_item_idx += 1
+
+            second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
+            if second_per_grid_ts:
+                video_second_per_grid_t = second_per_grid_ts[item_idx]
+            else:
+                video_second_per_grid_t = 1.0
+
+            return MRotaryEmbedding.omni_get_updates_use_audio_in_video(
+                thinker_config=thinker_config,
+                audio_len=audio_num_features,
+                video_grid_thw=video_grid_thw,
+                video_second_per_grid_t=video_second_per_grid_t,
+            )
+
+        video_replacement_fn = (
+            get_replacement_qwen2_use_audio_in_video
+            if use_audio_in_video
+            else partial(get_replacement_qwen2_vision, modality="video")
+        )
+
+        return [
+            PromptReplacement(
+                modality="audio",
+                target=audio_token,
+                replacement=get_replacement_qwen2_audio,
+            ),
+            PromptReplacement(
+                modality="image",
+                target=image_token,
+                replacement=partial(get_replacement_qwen2_vision, modality="image"),
+            ),
+            PromptReplacement(
+                modality="video",
+                target=video_token,
+                replacement=video_replacement_fn,
+            ),
+        ]
+
+    def _apply_hf_processor_main(
+        self,
+        prompt: Union[str, list[int]],
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+        *,
+        enable_hf_prompt_update: bool,
+    ) -> tuple[list[int], BatchFeature, bool]:
+        """
+        Qwen2.5-Omni reimplements this function to handle text only.
+        """
+        if isinstance(prompt, str):
+            if enable_hf_prompt_update:
+                return self._apply_hf_processor_text_mm(
+                    prompt_text=prompt,
+                    mm_items=mm_items,
+                    hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+                    tokenization_kwargs=tokenization_kwargs,
+                )
+            tokenizer = self.info.get_tokenizer()
+            prompt_ids = encode_tokens(tokenizer, prompt)
+        else:
+            prompt_ids = self._apply_hf_processor_tokens_only(prompt)
+
+        mm_processed_data = self._apply_hf_processor_mm_only(
+            mm_items=mm_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+            tokenization_kwargs=tokenization_kwargs,
+        )
+
+        return prompt_ids, mm_processed_data, False
+
+    def _apply_hf_processor_mm_only(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        """
+        Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
+        """
+        mm_counts = mm_items.get_all_counts()
+
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
+        if use_audio_in_video and "video" in mm_counts:
+            assert "audio" in mm_counts
+            mm_counts["audio"] -= mm_counts["video"]
+
+        _, mm_processed_data, _ = self._apply_hf_processor_text_mm(
+            prompt_text=self.dummy_inputs.get_dummy_text(mm_counts),
+            mm_items=mm_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+            tokenization_kwargs=tokenization_kwargs,
+        )
+
+        return mm_processed_data
+
+    def _validate_mm_placeholders(
+        self,
+        mm_placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
+        mm_item_counts: Mapping[str, int],
+        use_audio_in_video: bool = False,
+    ) -> None:
+        if use_audio_in_video:
+            mm_item_counts = copy(mm_item_counts)
+            if "video" in mm_item_counts:
+                assert "audio" in mm_item_counts
+                mm_item_counts["audio"] -= mm_item_counts["video"]
+        super()._validate_mm_placeholders(mm_placeholders, mm_item_counts)
+
+
+class Qwen2_5OmniConditionalGenerationMixin:
+    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}. 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]:
+        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)
+        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(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,
+            feature_attention_mask=feature_attention_mask,
+        )
+
+    def _parse_and_validate_image_input(
+        self,
+        **kwargs: dict[str, Any],
+    ) -> Optional[Qwen2_5_VLImageInputs]:
+        pixel_values = kwargs.pop("pixel_values", None)
+        image_embeds = kwargs.pop("image_embeds", None)
+        image_grid_thw = kwargs.pop("image_grid_thw", None)
+
+        if pixel_values is None and image_embeds is None:
+            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")
+
+            if not isinstance(pixel_values, (torch.Tensor, list)):
+                raise ValueError(f"Incorrect type of image pixel values. Got type: {type(pixel_values)}")
+
+            return Qwen2_5_VLImagePixelInputs(
+                type="pixel_values",
+                pixel_values=pixel_values,
+                image_grid_thw=image_grid_thw,
+            )
+
+        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")
+
+            if not isinstance(image_embeds, torch.Tensor):
+                raise ValueError(f"Incorrect type of image embeddings. Got type: {type(image_embeds)}")
+            return Qwen2_5_VLImageEmbeddingInputs(
+                type="image_embeds",
+                image_embeds=image_embeds,
+                image_grid_thw=image_grid_thw,
+            )
+
+    def _parse_and_validate_video_input(
+        self,
+        **kwargs: dict[str, Any],
+    ) -> Optional[Qwen2_5_VLVideoInputs]:
+        pixel_values_videos = kwargs.pop("pixel_values_videos", None)
+        video_embeds = kwargs.pop("video_embeds", None)
+        video_grid_thw = kwargs.pop("video_grid_thw", None)
+
+        if pixel_values_videos is None and video_embeds is None:
+            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")
+
+            return Qwen2_5_VLVideoPixelInputs(
+                type="pixel_values_videos",
+                pixel_values_videos=pixel_values_videos,
+                video_grid_thw=video_grid_thw,
+            )
+
+        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")
+
+            if not isinstance(video_embeds, torch.Tensor):
+                raise ValueError(f"Incorrect type of video embeddings. Got type: {type(video_embeds)}")
+            return Qwen2_5_VLVideoEmbeddingInputs(
+                type="video_embeds",
+                video_embeds=video_embeds,
+                video_grid_thw=video_grid_thw,
+            )
+
+    def _process_audio_input(
+        self,
+        audio_input: Qwen2AudioInputs,
+        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
+            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_outputs = self.audio_tower(
+            input_features.to(self.audio_tower.dtype),
+            feature_lens=audio_feature_lengths,
+            aftercnn_lens=audio_feat_lengths,
+        )
+        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, ...]:
+        if image_input["type"] == "image_embeds":
+            return image_input["image_embeds"].type(self.visual.dtype)
+
+        grid_thw = image_input["image_grid_thw"]
+        assert grid_thw.ndim == 2
+
+        pixel_values = image_input["pixel_values"].type(self.visual.dtype)
+        image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
+        # Split concatenated embeddings for each image item.
+        merge_size = self.visual.spatial_merge_size
+        sizes = grid_thw.prod(-1) // merge_size // merge_size
+
+        return image_embeds.split(sizes.tolist())
+
+    def _process_video_input(
+        self,
+        video_input: Qwen2_5_VLVideoInputs,
+        video_hashes: list[str] = None,
+        cached_video_embeds: torch.Tensor = None,
+    ) -> torch.Tensor:
+        if video_input["type"] == "video_embeds":
+            return video_input["video_embeds"].type(self.visual.dtype)
+
+        grid_thw = video_input["video_grid_thw"]
+        assert grid_thw.ndim == 2
+
+        pixel_values_videos = video_input["pixel_values_videos"].type(self.visual.dtype)
+        video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)
+        # Split concatenated embeddings for each video item.
+        merge_size = self.visual.spatial_merge_size
+        sizes = grid_thw.prod(-1) // merge_size // merge_size
+
+        return video_embeds.split(sizes.tolist())
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    info=Qwen2_5OmniThinkerProcessingInfo,
+    dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
+)
+class Qwen2_5OmniThinkerForConditionalGeneration(
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "thinker.lm_head.": "language_model.lm_head.",
+            "thinker.model.": "language_model.model.",
+            "thinker.": "",
+        }
+    )
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
+        if modality.startswith("image"):
+            return "<|vision_start|><|IMAGE|><|vision_end|>"
+        if modality.startswith("video"):
+            return "<|vision_start|><|VIDEO|><|vision_end|>"
+        if modality.startswith("audio"):
+            return f"Audio {i}: <|audio_bos|><|AUDIO|><|audio_eos|>"
+
+        raise ValueError("Only image, video or audio modality is supported")
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        thinker_config: Qwen2_5OmniThinkerConfig = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+        self.config = thinker_config
+        self.multimodal_config = multimodal_config
+
+        # force "use_flash_attention_2=True" to audio tower to align
+        # the results.
+        if flash_attn is not None:
+            audio_config = thinker_config.audio_config
+            audio_config._attn_implementation_autoset = True
+            audio_config._attn_implementation = "flash_attention_2"
+        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."
+            )
+
+        self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
+        self.visual = Qwen2_5_VisionTransformer(
+            vision_config=thinker_config.vision_config,
+            norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "visual"),
+        )
+        self.quant_config = quant_config
+        self.language_model = init_vllm_registered_model(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "language_model"),
+            hf_config=thinker_config.text_config,
+            architectures=["Qwen2ForCausalLM"],
+        )
+
+        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 = {}
+
+        # 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)
+        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 []
+
+        # The result multimodal_embeddings is tuple of tensors, with each
+        # tensor correspoending to a multimodal data item (image or video).
+        multimodal_embeddings: tuple[torch.Tensor, ...] = ()
+
+        # NOTE: It is important to iterate over the keys in this dictionary
+        # to preserve the order of the modalities.
+        for modality in mm_input_by_modality:
+            multimodal_input = mm_input_by_modality[modality]
+            if modality == "image":
+                vision_embeddings = self._process_image_input(multimodal_input)
+                multimodal_embeddings += vision_embeddings
+            if modality == "video":
+                video_embeddings = self._process_video_input(multimodal_input)
+                multimodal_embeddings += video_embeddings
+            if modality == "audio":
+                audio_embeddings = self._process_audio_input(multimodal_input)
+                multimodal_embeddings += audio_embeddings
+        return multimodal_embeddings
+
+    def get_input_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
+    ) -> 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(
+                input_ids,
+                inputs_embeds,
+                multimodal_embeddings,
+                [
+                    self.config.image_token_index,
+                    self.config.video_token_index,
+                    self.config.audio_token_index,
+                ],
+            )
+        return inputs_embeds
+
+    def get_multimodal_embeddings_v0(self, **kwargs: object) -> Optional[NestedTensors]:
+        audio_input = self._parse_and_validate_audio_input(**kwargs)
+        image_input = self._parse_and_validate_image_input(**kwargs)
+        video_input = self._parse_and_validate_video_input(**kwargs)
+
+        if audio_input is None and image_input is None and video_input is None:
+            return None
+
+        multimodal_embeddings: list[tuple[NestedTensors, str]] = []
+
+        if audio_input is not None:
+            audio_embeds = self._process_audio_input(audio_input)
+            multimodal_embeddings.append((audio_embeds, "audio"))
+        if image_input is not None:
+            image_embeds = self._process_image_input(image_input)
+            multimodal_embeddings.append((image_embeds, "image"))
+        if video_input is not None:
+            video_embeds = self._process_video_input(video_input)
+            multimodal_embeddings.append((video_embeds, "video"))
+        return multimodal_embeddings
+
+    def get_input_embeddings_v0(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: Optional[NestedTensors] = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+
+        for embeddings, modality in multimodal_embeddings:
+            if modality == "audio":
+                placeholder_token_id = self.config.audio_token_index
+            if modality == "image":
+                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)
+        return inputs_embeds
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs: object,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+
+        # NOTE: In v1, inputs_embeds is always generated at model runner, this
+        # condition is for v0 compatibility.
+        elif inputs_embeds is None:
+            multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)
+            inputs_embeds = self.get_input_embeddings(input_ids, multimodal_embeddings)
+            text_inputs_embeds = self.get_input_embeddings(
+                input_ids,
+                (
+                    [(torch.zeros_like(embeddings), modality) for embeddings, modality in multimodal_embeddings]
+                    if multimodal_embeddings is not None
+                    else None
+                ),
+            )
+            input_ids = None
+        else:
+            text_inputs_embeds = inputs_embeds
+
+        hidden_states = self.language_model.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
+        )
+        return text_inputs_embeds, hidden_states.unsqueeze(0)  # (1, S, D)
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:
+        return self.language_model.compute_logits(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=["talker.", "token2wav."],
+        )
+        loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        # Log load summary
+        try:
+            total_bytes = 0
+            for name, param in self.named_parameters():
+                if param is not None and param.data is not None:
+                    total_bytes += param.data.numel() * param.data.element_size()
+            device = next(self.parameters()).device
+            logger.info(
+                "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
+                self.__class__.__name__,
+                True,
+                total_bytes / (1024**2),
+                str(device),
+            )
+        except Exception:
+            pass
+        return loaded_weights

vllm_omni/model_executor/models/qwen2_old.py

@@ -0,0 +1,520 @@
+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
+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.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.models.interfaces import SupportsLoRA, SupportsPP
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import PoolerOutput, SamplerOutput
+from vllm.v1.pool.metadata import PoolingMetadata
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
+
+logger = init_logger(__name__)
+
+
+class Qwen2MLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.")
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Qwen2Attention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        max_position: int = 4096 * 32,
+        rope_theta: float = 10000,
+        head_dim: Optional[int] = None,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        rope_scaling: Optional[tuple] = None,
+        prefix: str = "",
+        attn_type: str = AttentionType.DECODER,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim or (hidden_size // self.total_num_heads)
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=max_position,
+            base=self.rope_theta,
+            rope_scaling=rope_scaling,
+        )
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+            attn_type=attn_type,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Qwen2DecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: Qwen2Config,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        # Requires transformers > 4.32.0
+        rope_theta = getattr(config, "rope_theta", 1000000)
+        rope_scaling = getattr(config, "rope_scaling", None)
+
+        # By default, Qwen2 uses causal attention as it is a decoder-only model.
+        # You can override the HF config with `is_causal=False` to enable
+        # bidirectional attention, which is used in some embedding models
+        # (e.g. Alibaba-NLP/gte-Qwen2-7B-instruct)
+        if getattr(config, "is_causal", True):
+            attn_type = AttentionType.DECODER
+        else:
+            attn_type = AttentionType.ENCODER_ONLY
+
+        self.self_attn = Qwen2Attention(
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            max_position=config.max_position_embeddings,
+            num_kv_heads=config.num_key_value_heads,
+            rope_theta=rope_theta,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            rope_scaling=rope_scaling,
+            head_dim=getattr(config, "head_dim", None),
+            prefix=f"{prefix}.self_attn",
+            attn_type=attn_type,
+        )
+        self.mlp = Qwen2MLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            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)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
+        # otherwise (seq_len, ).
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+    }
+)
+class Qwen2Model(nn.Module):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+        decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer,
+    ):
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        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"):
+            raise ValueError(
+                "Sliding window for some but all layers is not "
+                "supported. This model uses sliding window "
+                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):
+            self.embed_tokens = VocabParallelEmbedding(
+                config.vocab_size,
+                getattr(config, "embedding_size", config.hidden_size),
+                quant_config=quant_config,
+                prefix=f"{prefix}.embed_tokens",
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        # Use the provided decoder layer type or default to Qwen2DecoderLayer
+        decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: decoder_layer_type(
+                config=config,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=prefix,
+            ),
+            prefix=f"{prefix}.layers",
+        )
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.get_input_embeddings(input_ids)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+        for layer in self.layers[self.start_layer : self.end_layer]:
+            hidden_states, residual = layer(
+                positions,
+                hidden_states,
+                residual,
+            )
+        if not get_pp_group().is_last_rank:
+            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]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        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)):
+                # 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]
+                weight_loader(param, loaded_weight)
+                loaded_params.add(scale_name)
+                continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                # Remapping the name of FP8 kv-scale.
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "gate_up_proj": [
+            "gate_proj",
+            "up_proj",
+        ],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        lora_config = vllm_config.lora_config
+
+        self.config = config
+        self.lora_config = lora_config
+
+        self.quant_config = quant_config
+        self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+
+        if get_pp_group().is_last_rank:
+            if config.tie_word_embeddings:
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.sampler = Sampler()
+
+        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)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        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)
+        return hidden_states
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        return logits
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[SamplerOutput]:
+        next_tokens = self.sampler(logits, sampling_metadata)
+        return next_tokens
+
+    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),
+        )
+        return loader.load_weights(weights)
+
+
+class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "gate_up_proj": [
+            "gate_proj",
+            "up_proj",
+        ],
+    }
+
+    hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        lora_config = vllm_config.lora_config
+        pooler_config = vllm_config.model_config.pooler_config
+
+        self.config = config
+        self.lora_config = lora_config
+
+        self.quant_config = quant_config
+        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
+        if pooler_config.pooling_type is None:
+            logger.warning(
+                "This embedding model will default to last-token pooling in "
+                "an upcoming version. To avoid breaking changes, you should "
+                'pass `--override-pooler-config \'{"pooling_type": "MEAN"}\'`'
+                " explicitly."
+            )
+
+        self._pooler = Pooler.from_config_with_defaults(
+            pooler_config, pooling_type=PoolingType.MEAN, normalize=True, softmax=False
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> torch.Tensor:
+        return self.model(input_ids, positions, intermediate_tensors)
+
+    def pooler(
+        self,
+        hidden_states: torch.Tensor,
+        pooling_metadata: PoolingMetadata,
+    ) -> Optional[PoolerOutput]:
+        return self._pooler(hidden_states, pooling_metadata)
+
+    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."))
+        self.model.load_weights(weights)

vllm_omni/model_executor/models/registry.py

@@ -0,0 +1,50 @@
+from vllm.model_executor.models.registry import _VLLM_MODELS, _LazyRegisteredModel, _ModelRegistry
+
+_OMNI_MODELS = {
+    "Qwen2_5OmniForConditionalGeneration": (
+        "qwen2_5_omni",
+        "Qwen2_5OmniForConditionalGeneration",
+    ),
+    "Qwen2_5OmniThinkerModel": (
+        "qwen2_5_omni_thinker",
+        "Qwen2_5OmniThinkerForConditionalGeneration",
+    ),  # noqa: E501
+    "Qwen2_5OmniTalkerModel": (
+        "qwen2_5_omni_talker",
+        "Qwen2_5OmniTalkerForConditionalGeneration",
+    ),  # noqa: E501
+    "Qwen2_5OmniToken2WavModel": (
+        "qwen2_5_omni_token2wav",
+        "Qwen2_5OmniToken2WavForConditionalGenerationVLLM",
+    ),
+    "Qwen2_5OmniToken2WavDiTModel": (
+        "qwen2_5_omni_token2wav",
+        "Qwen2_5OmniToken2WavModel",
+    ),
+    "Qwen2ForCausalLM_old": ("qwen2_old", "Qwen2ForCausalLM"),  # need to discuss
+}
+
+_VLLM_OMNI_MODELS = {
+    **_VLLM_MODELS,
+    **_OMNI_MODELS,
+}
+
+
+OmniModelRegistry = _ModelRegistry(
+    {
+        **{
+            model_arch: _LazyRegisteredModel(
+                module_name=f"vllm.model_executor.models.{mod_relname}",
+                class_name=cls_name,
+            )
+            for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()
+        },
+        **{
+            model_arch: _LazyRegisteredModel(
+                module_name=f"vllm_omni.model_executor.models.{mod_relname}",
+                class_name=cls_name,
+            )
+            for model_arch, (mod_relname, cls_name) in _OMNI_MODELS.items()
+        },
+    }
+)

vllm_omni/model_executor/models/utils.py

@@ -0,0 +1,8 @@
+from vllm.model_executor.models.utils import maybe_prefix
+
+
+def add_prefix_to_loaded_weights(weights: set[str], prefix: str) -> set[str]:
+    """
+    Add a prefix to the names of the loaded weights.
+    """
+    return {maybe_prefix(prefix, name) for name in weights}

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"]