diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index a67d038e1..7bc41512e 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -1,7 +1,6 @@
steps:
- - label: "Dummy pipeline"
- if: "false"
- agents:
- queue: cpu_queue_premerge_us_east_1
+ - label: "Simple Unit Test"
commands:
- - "echo 'This pipeline is intentionally a no-op.'"
+ - ".buildkite/scripts/simple_test.sh"
+ agents:
+ queue: "cpu_queue_premerge"
diff --git a/.buildkite/scripts/simple_test.sh b/.buildkite/scripts/simple_test.sh
new file mode 100755
index 000000000..6c406d095
--- /dev/null
+++ b/.buildkite/scripts/simple_test.sh
@@ -0,0 +1,52 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)"
+UV_BIN="uv"
+PYTHON_BOOTSTRAP="${PYTHON:-python3}"
+VENV_DIR="${VENV_DIR:-${REPO_ROOT}/.venv-simple-test}"
+TARGET_PY_VERSION="3.12"
+
+if ! command -v "${PYTHON_BOOTSTRAP}" >/dev/null 2>&1; then
+ PYTHON_BOOTSTRAP="python"
+fi
+
+cd "${REPO_ROOT}"
+
+# Ensure pip-installed scripts (including uv) are on PATH
+PYTHON_SCRIPTS="$(${PYTHON_BOOTSTRAP} -c 'import sysconfig; print(sysconfig.get_path("scripts"))')"
+USER_BASE="$(${PYTHON_BOOTSTRAP} -m site --user-base 2>/dev/null || true)"
+PATH="${PYTHON_SCRIPTS}:${PATH}"
+if [[ -n "${USER_BASE}" ]]; then
+ PATH="${USER_BASE}/bin:${PATH}"
+fi
+
+if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
+ "${PYTHON_BOOTSTRAP}" -m pip install --upgrade pip
+ "${PYTHON_BOOTSTRAP}" -m pip install uv
+ hash -r
+fi
+
+TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
+if [[ -z "${TARGET_PY_PATH}" ]]; then
+ "${UV_BIN}" python install "${TARGET_PY_VERSION}"
+ TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
+fi
+
+if [[ -z "${TARGET_PY_PATH}" ]]; then
+ echo "Failed to locate python ${TARGET_PY_VERSION}" >&2
+ exit 1
+fi
+
+if [[ ! -d "${VENV_DIR}" ]]; then
+ "${UV_BIN}" venv --python "${TARGET_PY_PATH}" "${VENV_DIR}"
+fi
+
+VENV_PYTHON="${VENV_DIR}/bin/python"
+[[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
+
+"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.11.0
+"${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
+"${VENV_PYTHON}" -m pytest -v -s tests/test_omni_llm.py
diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml
index ff0770175..ef02d5747 100644
--- a/.github/ISSUE_TEMPLATE/config.yml
+++ b/.github/ISSUE_TEMPLATE/config.yml
@@ -2,4 +2,4 @@ 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
+ about: Ask questions and discuss with other vLLM-Omni community members
diff --git a/.gitignore b/.gitignore
index 59a839811..e38e2dc74 100644
--- a/.gitignore
+++ b/.gitignore
@@ -175,7 +175,7 @@ tests/
# *.iws
# .idea/
-# vLLM-omni specific
+# vLLM-Omni specific
# Model files and checkpoints
checkpoints/
*.bin
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 486b2b10a..4f610059a 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -1,19 +1,3 @@
-# Contributing to vLLM
+# Contributing to vLLM-Omni
-## 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
-```
+You may find information about contributing to vLLM-Omni on [Contributing](docs/contributing/)
diff --git a/README.md b/README.md
index 52001e60e..6201c4b44 100644
--- a/README.md
+++ b/README.md
@@ -22,7 +22,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
## About
-[vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni extends its support for omni-modality model inference and serving:
+[vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:
- **Omni-modality**: Text, image, video, and audio data processing
- **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
@@ -34,98 +34,39 @@ Easy, fast, and cheap omni-modality model serving for everyone
-vLLM-omni is fast with:
+vLLM-Omni is fast with:
-- Seamless AR support for efficient KV Cache management with vLLM
-- Pipelined stage execution overlapping
-- Fully disaggregation based on omniConnector and dynamic resource allocation across stages
+- State-of-the-art AR support by leveraging efficient KV cache management from vLLM
+- Pipelined stage execution overlapping for high throughput performance
+- Fully disaggregation based on OmniConnector and dynamic resource allocation across stages
-vLLM-omni is flexible and easy to use with:
+vLLM-Omni is flexible and easy to use with:
+- Heterogeneous pipeline abstraction to manage complex model workflows
- Seamless integration with popular Hugging Face models
-- High-throughput serving with various decoding algorithms, including *parallel sampling*, *beam search*, and more
- Tensor, pipeline, data and expert parallelism support for distributed inference
- Streaming outputs
- OpenAI-compatible API server
-vLLM seamlessly supports most popular open-source models on HuggingFace, including:
+vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
-- Multi-modality generation models (e.g. Qwen-image)
-- Omni-modality models (e.g. Qwen-omni)
+- Omni-modality models (e.g. Qwen-Omni)
+- Multi-modality generation models (e.g. Qwen-Image)
## Getting Started
-Install vLLM with `pip` or [from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source):
+Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn more.
-```bash
-pip install vllm
-```
+Visit our documentation to learn more.
-## 🛠️ 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 .
-```
-
-### 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)
+- [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
+- [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
+- [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
## Contributing
We welcome and value any contributions and collaborations.
-Please check out [Contributing to vLLM-omni](https://docs.vllm.ai/en/latest/contributing/index.html) for how to get involved.
+Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.
## Weekly Meeting
diff --git a/docs/.nav.yml b/docs/.nav.yml
new file mode 100644
index 000000000..a9cd65d33
--- /dev/null
+++ b/docs/.nav.yml
@@ -0,0 +1,44 @@
+nav:
+ - Home: README.md
+ - User Guide:
+ - Getting Started:
+ - getting_started/quickstart.md
+ - getting_started/installation
+ - Examples:
+ - examples/README.md
+ - Offline Inference:
+ - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
+ - Online Serving:
+ - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
+ - Qwen2.5-Image: user_guide/examples/offline_inference/qwen_image.md
+ - General:
+ - usage/*
+ - Configuration:
+ - configuration/README.md
+ - configuration/*
+ - Models:
+ - models/supported_models.md
+ - Developer Guide:
+ - General:
+ - contributing/README.md
+ - glob: contributing/*
+ flatten_single_child_sections: true
+ - Model Implementation:
+ - contributing/model/README.md
+ - CI: contributing/ci
+ - Design Documents:
+ - design/index.md
+ - design/architecture_overview.md
+ - design/vllm_omni_design.md
+ - design/mrs_design.md
+ - design/api_design_doc.md
+ - Docs Guide: contributing/DOCS_GUIDE.md
+ - API Reference:
+ - api/README.md
+ - api/vllm_omni
+ - CLI Reference: cli
+ - Community:
+ - community/*
+ - Slack: https://slack.vllm.ai
+ - Blog: https://blog.vllm.ai
+ - Forum: https://discuss.vllm.ai
diff --git a/docs/README.md b/docs/README.md
new file mode 100644
index 000000000..27bf83f39
--- /dev/null
+++ b/docs/README.md
@@ -0,0 +1,61 @@
+---
+hide:
+ - navigation
+ - toc
+---
+
+# Welcome to vLLM-Omni
+
+
+
+
+
+
+
+Easy, fast, and cheap omni-modality model serving for everyone
+
+
+
+
+Star
+Watch
+Fork
+
+
+
+## About
+
+[vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:
+
+- **Omni-modality**: Text, image, video, and audio data processing
+- **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
+- **Heterogeneous outputs**: from traditional text generation to multimodal outputs
+
+
+
+
+
+
+vLLM-Omni is fast with:
+
+- State-of-the-art AR support by leveraging efficient KV cache management from vLLM
+- Pipelined stage execution overlapping for high throughput performance
+- Fully disaggregation based on OmniConnector and dynamic resource allocation across stages
+
+vLLM-Omni is flexible and easy to use with:
+
+- Heterogeneous pipeline abstraction to manage complex model workflows
+- Seamless integration with popular Hugging Face models
+- Tensor, pipeline, data and expert parallelism support for distributed inference
+- Streaming outputs
+- OpenAI-compatible API server
+
+vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
+
+- Omni-modality models (e.g. Qwen2.5-Omni, Qwen3-Omni)
+- Multi-modality generation models (e.g. Qwen-Image)
+
+For more information, checkout the following:
+- [Placeholder for vllm-omni blogpost]()
diff --git a/docs/api/README.md b/docs/api/README.md
new file mode 100644
index 000000000..100e17ad4
--- /dev/null
+++ b/docs/api/README.md
@@ -0,0 +1,92 @@
+# Summary
+
+## Entry Points
+
+Main entry points for vLLM-Omni inference and serving.
+
+- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
+- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniStageLLM][]
+- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
+- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
+- [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
+- [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
+- [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][]
+- [vllm_omni.entrypoints.omni.Omni][]
+- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
+- [vllm_omni.entrypoints.omni_llm.OmniStageLLM][]
+- [vllm_omni.entrypoints.omni_stage.OmniStage][]
+- [vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat][]
+
+## Inputs
+
+Input data structures for multi-modal inputs.
+
+- [vllm_omni.inputs.data.OmniEmbedsPrompt][]
+- [vllm_omni.inputs.data.OmniTokenInputs][]
+- [vllm_omni.inputs.data.OmniTokensPrompt][]
+- [vllm_omni.inputs.parse.parse_singleton_prompt_omni][]
+- [vllm_omni.inputs.preprocess.OmniInputPreprocessor][]
+
+## Engine
+
+Engine classes for offline and online inference.
+
+- [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
+- [vllm_omni.engine.AdditionalInformationEntry][]
+- [vllm_omni.engine.AdditionalInformationPayload][]
+- [vllm_omni.engine.OmniEngineCoreRequest][]
+- [vllm_omni.engine.PromptEmbedsPayload][]
+- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
+- [vllm_omni.engine.arg_utils.OmniEngineArgs][]
+- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
+- [vllm_omni.engine.output_processor.OmniRequestState][]
+- [vllm_omni.engine.processor.OmniProcessor][]
+
+## Core
+
+Core scheduling and caching components.
+
+- [vllm_omni.core.dit_cache_manager.DiTCacheManager][]
+- [vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler][]
+- [vllm_omni.core.sched.output.OmniNewRequestData][]
+- [vllm_omni.core.sched.scheduler.OmniScheduler][]
+
+## Model Executor
+
+Model execution components.
+
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.OmniOutput][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniAudioFeatureInputs][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniConditionalGenerationMixin][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniThinkerDummyInputsBuilder][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniThinkerForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniThinkerMultiModalDataParser][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniThinkerMultiModalProcessor][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniThinkerProcessingInfo][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavBigVGANModel][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavDiTModel][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavForConditionalGenerationVLLM][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavModel][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2EmbeddingModel][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2ForCausalLM][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2Model][]
+
+## Configuration
+
+Configuration classes.
+
+- [vllm_omni.config.model.OmniModelConfig][]
+
+## Workers
+
+Worker classes and model runners for distributed inference.
+
+- [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][]
+- [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][]
+- [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
+- [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
+- [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
+- [vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
+- [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
diff --git a/docs/api/overview.md b/docs/api/overview.md
deleted file mode 100644
index 8a33a9095..000000000
--- a/docs/api/overview.md
+++ /dev/null
@@ -1,68 +0,0 @@
-# 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][]
diff --git a/docs/cli/README.md b/docs/cli/README.md
new file mode 100644
index 000000000..8069110b3
--- /dev/null
+++ b/docs/cli/README.md
@@ -0,0 +1,24 @@
+# vLLM-Omni CLI Guide
+
+The CLI for vLLM-Omni inherits from vllm with some additional arguments.
+
+## serve
+
+Starts the vLLM-Omni OpenAI Compatible API server.
+
+Start with a model:
+
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni
+```
+
+Specify the port:
+
+```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 --stage-configs-path /path/to/stage_configs_file
+```
diff --git a/docs/community/contact_us.md b/docs/community/contact_us.md
new file mode 100644
index 000000000..8edc2b957
--- /dev/null
+++ b/docs/community/contact_us.md
@@ -0,0 +1,5 @@
+# Contact Us
+
+- For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
+- For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/)
+- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm/security/advisories) feature
diff --git a/docs/community/meetups.md b/docs/community/meetups.md
new file mode 100644
index 000000000..3374fe711
--- /dev/null
+++ b/docs/community/meetups.md
@@ -0,0 +1 @@
+# Meetups
diff --git a/docs/configuration/README.md b/docs/configuration/README.md
new file mode 100644
index 000000000..02f164a03
--- /dev/null
+++ b/docs/configuration/README.md
@@ -0,0 +1,11 @@
+# Configuration Options
+
+This section lists the most common options for running vLLM-Omni.
+
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.11.0/configuration/index.html)
+
+Currently, the main options are maintained by stage configs for each model.
+
+For specific example, please refer to [Qwen2.5-omni stage config](stage_configs/qwen2_5_omni.yaml)
+
+For introduction, please check [Introduction for stage config](./stage_configs.md)
diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md
new file mode 100644
index 000000000..2ed46e93d
--- /dev/null
+++ b/docs/configuration/stage_configs.md
@@ -0,0 +1,269 @@
+# Stage configs for vLLM-Omni
+
+In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion.
+
+Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions:
+
+- Claim partition of stages and their corresponding class implementation in `model_executor/models`.
+- The disaggregated configuration for each stage and the communication topology among them.
+- Engine arguments for each engine within the stage.
+- Input and output dependencies for each stage.
+- Default input parameters.
+
+If users want to modify some part of it. The custom stage_configs file can be input as input argument in both online and offline. Just like examples below:
+
+For offline (Assume necessary dependencies have ben imported):
+```python
+model_name = "Qwen/Qwen2.5-Omni-7B"
+omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+```
+
+For online serving:
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+```
+!!! important
+ We are actively iterating on the definition of stage configs, and we welcome all feedbacks from both community users and developers to help us shape the development!
+
+Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
+```python
+# stage config for running qwen2.5-omni with architecture of OmniLLM.
+stage_args:
+ - stage_id: 0 # mark the unique id for each stage
+ runtime: # The disaggregated configuration
+ process: true # Run this stage in a separate process
+ devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+ max_batch_size: 1 # the batch_size for offline inference
+ engine_args: # Engine arguments for a certain engine
+ model_stage: thinker
+ model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
+ worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker # The specific worker used
+ scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler # The specific scehduler used
+ gpu_memory_utilization: 0.8 # The gpu memory allocation for the stage within a single chip
+ enforce_eager: true # Now we only support eager mode
+ trust_remote_code: true # Needed by huggingface config parsing
+ engine_output_type: latent # It claims that the stage will input latent hiddenstates besides token ids
+ enable_prefix_caching: false # For request with hiddenstates output, the prefix caching is not supported now
+ is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmniLLM will use its tokenizer as default
+ final_output: true # If the stage has output as part of final outputs. If it is false, which means that the stage only works as a intermediate role.
+ final_output_type: text # What is the final output type. It can be text and audio now.
+ default_sampling_params: # sampling parameters for the stage. Their meaning aligns with vLLM.
+ temperature: 0.0
+ top_p: 1.0
+ top_k: -1
+ max_tokens: 2048
+ seed: 42
+ detokenize: True
+ repetition_penalty: 1.1
+ - stage_id: 1
+ runtime:
+ process: true
+ devices: "1"
+ max_batch_size: 3
+ engine_args:
+ model_stage: talker
+ model_arch: Qwen2_5OmniForConditionalGeneration
+ worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+ scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+ gpu_memory_utilization: 0.8
+ enforce_eager: true
+ trust_remote_code: true
+ enable_prefix_caching: false
+ engine_output_type: latent
+ engine_input_source: [0]
+ custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
+ default_sampling_params:
+ temperature: 0.9
+ top_p: 0.8
+ top_k: 40
+ max_tokens: 2048
+ seed: 42
+ detokenize: True
+ repetition_penalty: 1.05
+ stop_token_ids: [8294]
+ - stage_id: 2
+ runtime:
+ process: true
+ devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
+ max_batch_size: 1
+ engine_args:
+ model_stage: code2wav
+ model_arch: Qwen2_5OmniForConditionalGeneration
+ worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
+ scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
+ gpu_memory_utilization: 0.15
+ enforce_eager: true
+ trust_remote_code: true
+ enable_prefix_caching: false
+ engine_output_type: audio
+ engine_input_source: [1]
+ final_output: true
+ final_output_type: audio
+ default_sampling_params:
+ temperature: 0.0
+ top_p: 1.0
+ top_k: -1
+ max_tokens: 2048
+ seed: 42
+ detokenize: True
+ repetition_penalty: 1.1
+
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
+ enabled: true
+ defaults:
+ window_size: -1 # Simplified: trigger downstream only after full upstream completion
+ max_inflight: 1 # Simplified: process serially within each stage
+ edges:
+ - from: 0 # thinker → talker: trigger only after receiving full input (-1)
+ to: 1
+ window_size: -1
+ - from: 1 # talker → code2wav: trigger only after receiving full input (-1)
+ to: 2
+ window_size: -1
+
+```
+
+## Stage Configuration Arguments
+
+Each stage in the `stage_args` list contains the following configuration options:
+
+### `stage_id`
+
+A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).
+
+### `runtime`
+
+Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.
+
+#### `runtime.process`
+
+Whether to run this stage in a separate process. When set to `true`, the stage will be executed in an isolated process, enabling better resource isolation and parallel execution across different stages. This is essential for multi-GPU deployments where different stages run on different devices.
+
+Default: `true`
+
+#### `runtime.devices`
+
+Visible devices for this stage, specified as a string. This controls which GPU devices are available to the stage process, similar to setting `CUDA_VISIBLE_DEVICES` or using `torch.cuda.set_device()`. For example, `"0"` uses GPU 0, `"1"` uses GPU 1, and `"0,1"` makes both GPUs 0 and 1 visible.
+
+Default: `"0"`
+
+#### `runtime.max_batch_size`
+
+The maximum batch size for offline inference in this stage. This limits how many sequences can be processed together in a single batch during offline inference operations.
+
+Default: `1`
+
+### `engine_args`
+
+Engine arguments for configuring the LLM engine, diffusion engine, or other engine types used by this stage.
+
+#### `engine_args.model_stage`
+
+The name identifier for this model stage within the multi-stage architecture. This is used internally to distinguish different stages of the same model (e.g., "thinker", "talker", "code2wav" in Qwen2.5-Omni).
+
+#### `engine_args.model_arch`
+
+The model architecture class name that is registered in `model_executor/models/registry.py`. This specifies which model implementation to use for this stage. The class must be registered in the model registry for vLLM-Omni to locate and instantiate it.
+
+#### `engine_args.worker_cls`
+
+The specific worker class to use for this stage. This determines how the model computations are executed. Examples include `vllm_omni.worker.gpu_ar_worker.GPUARWorker` for autoregressive stages and `vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker` for diffusion-based stages.
+
+#### `engine_args.scheduler_cls`
+
+The scheduler class to use for this stage. The scheduler manages request queuing, batching, and execution order. Examples include `vllm_omni.core.sched.scheduler.OmniScheduler` for standard stages and `vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler` for diffusion stages.
+
+#### `engine_args.gpu_memory_utilization`
+
+The fraction of GPU memory to allocate for this stage within a single GPU chip. This is a value between 0.0 and 1.0, where 0.8 means 80% of the GPU memory will be used by this stage. This allows fine-grained control over memory allocation when multiple stages share the same GPU or when reserving memory for other operations.
+
+Default: `0.8`
+
+#### `engine_args.enforce_eager`
+
+Whether to enforce eager execution mode. When set to `true`, the engine will run in eager mode without using CUDA graphs or other compilation optimizations. Currently, vLLM-Omni only supports eager mode.
+
+Default: `true`
+
+#### `engine_args.trust_remote_code`
+
+Whether to trust remote code when loading models from Hugging Face. This is required for models that use custom code in their configuration files. Set to `true` when loading models that require custom model implementations.
+
+Default: `true`
+
+#### `engine_args.engine_output_type`
+
+Specifies the type of output produced by this stage's engine. This determines what kind of data flows to downstream stages. Possible values include `latent` (hidden states), `text` (tokenized text), and `audio` (audio waveforms). When set to `latent`, the stage outputs latent hidden states in addition to token IDs, which are consumed by downstream stages.
+
+Default: `latent`
+
+#### `engine_args.enable_prefix_caching`
+
+Whether to enable prefix caching for this stage. Prefix caching can improve performance by caching KV cache for common prompt prefixes. However, for requests that output hidden states (when `engine_output_type` is `latent`), prefix caching is not currently supported and should be set to `false`.
+
+Default: `false`
+
+### `is_comprehension`
+
+Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmniLLM` will use its tokenizer as the default tokenizer for the entire pipeline.
+
+Default: `true`
+
+### `final_output`
+
+Whether this stage produces output that is part of the final outputs returned to the user. When set to `false`, the stage only works as an intermediate stage, processing data that flows to downstream stages but not contributing directly to the final response.
+
+Default: `true`
+
+### `final_output_type`
+
+The type of final output produced by this stage. This specifies what format the output will be in when returned to the user. Currently supported values are `text` (for text generation) and `audio` (for audio generation).
+
+Default: `text`
+
+### `default_sampling_params`
+
+Default sampling parameters for this stage. These parameters control the generation behavior and align with vLLM's sampling parameter semantics. These defaults are used when no explicit sampling parameters are provided in the request.
+
+#### `default_sampling_params.temperature`
+
+Sampling temperature for controlling randomness. Lower values (e.g., 0.0) make the output more deterministic and focused, while higher values increase randomness.
+
+Default: `0.0`
+
+#### `default_sampling_params.top_p`
+
+Nucleus sampling parameter. Only tokens with cumulative probability mass up to `top_p` are considered. This helps filter out low-probability tokens while maintaining diversity.
+
+Default: `1.0`
+
+#### `default_sampling_params.top_k`
+
+Top-k sampling parameter. Only the top `k` most likely tokens are considered. Set to `-1` to disable top-k filtering and consider all tokens.
+
+Default: `-1`
+
+#### `default_sampling_params.max_tokens`
+
+Maximum number of tokens to generate in this stage. This limits the length of the output sequence.
+
+Default: `2048`
+
+#### `default_sampling_params.seed`
+
+Random seed for reproducible generation. When set, the random number generator will be initialized with this seed to ensure consistent outputs across runs.
+
+Default: `42`
+
+#### `default_sampling_params.detokenize`
+
+Whether to detokenize the output tokens into text. When set to `true`, token IDs are converted back to readable text strings.
+
+Default: `True`
+
+#### `default_sampling_params.repetition_penalty`
+
+Penalty applied to tokens that have already appeared in the generated sequence. Values greater than 1.0 discourage repetition, while values less than 1.0 encourage it. A value of 1.0 applies no penalty.
+
+Default: `1.1`
diff --git a/docs/configuration/stage_configs/qwen2_5_omni.yaml b/docs/configuration/stage_configs/qwen2_5_omni.yaml
new file mode 100644
index 000000000..c378a5b48
--- /dev/null
+++ b/docs/configuration/stage_configs/qwen2_5_omni.yaml
@@ -0,0 +1,94 @@
+# stage config for running qwen2.5-omni with architecture of OmniLLM.
+stage_args:
+ - stage_id: 0
+ runtime:
+ process: true # Run this stage in a separate process
+ devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+ max_batch_size: 1
+ engine_args:
+ model_stage: thinker
+ model_arch: Qwen2_5OmniForConditionalGeneration
+ worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+ scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+ gpu_memory_utilization: 0.8
+ enforce_eager: true # Now we only support eager mode
+ trust_remote_code: true
+ engine_output_type: latent
+ enable_prefix_caching: false
+ is_comprehension: true
+ final_output: true
+ final_output_type: text
+ default_sampling_params:
+ temperature: 0.0
+ top_p: 1.0
+ top_k: -1
+ max_tokens: 2048
+ seed: 42
+ detokenize: True
+ repetition_penalty: 1.1
+ - stage_id: 1
+ runtime:
+ process: true
+ devices: "1"
+ max_batch_size: 1
+ engine_args:
+ model_stage: talker
+ model_arch: Qwen2_5OmniForConditionalGeneration
+ worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+ scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+ gpu_memory_utilization: 0.8
+ enforce_eager: true
+ trust_remote_code: true
+ enable_prefix_caching: false
+ engine_output_type: latent
+ engine_input_source: [0]
+ custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
+ default_sampling_params:
+ temperature: 0.9
+ top_p: 0.8
+ top_k: 40
+ max_tokens: 2048
+ seed: 42
+ detokenize: True
+ repetition_penalty: 1.05
+ stop_token_ids: [8294]
+ - stage_id: 2
+ runtime:
+ process: true
+ devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
+ max_batch_size: 1
+ engine_args:
+ model_stage: code2wav
+ model_arch: Qwen2_5OmniForConditionalGeneration
+ worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
+ scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
+ gpu_memory_utilization: 0.15
+ enforce_eager: true
+ trust_remote_code: true
+ enable_prefix_caching: false
+ engine_output_type: audio
+ engine_input_source: [1]
+ final_output: true
+ final_output_type: audio
+ default_sampling_params:
+ temperature: 0.0
+ top_p: 1.0
+ top_k: -1
+ max_tokens: 2048
+ seed: 42
+ detokenize: True
+ repetition_penalty: 1.1
+
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
+ enabled: true
+ defaults:
+ window_size: -1 # Simplified: trigger downstream only after full upstream completion
+ max_inflight: 1 # Simplified: process serially within each stage
+ edges:
+ - from: 0 # thinker → talker: trigger only after receiving full input (-1)
+ to: 1
+ window_size: -1
+ - from: 1 # talker → code2wav: trigger only after receiving full input (-1)
+ to: 2
+ window_size: -1
diff --git a/docs/DOCS_GUIDE.md b/docs/contributing/DOCS_GUIDE.md
similarity index 95%
rename from docs/DOCS_GUIDE.md
rename to docs/contributing/DOCS_GUIDE.md
index ccba0663f..dcbe2d7b6 100644
--- a/docs/DOCS_GUIDE.md
+++ b/docs/contributing/DOCS_GUIDE.md
@@ -1,6 +1,6 @@
# Documentation Build Guide
-This directory contains the source files for the vLLM-omni documentation.
+This directory contains the source files for the vLLM-Omni documentation.
## Building Documentation Locally
@@ -38,8 +38,8 @@ The documentation automatically extracts docstrings from the code using mkdocstr
Example docstring:
```python
-class OmniLLM:
- """Main entry point for vLLM-omni inference.
+class Omni:
+ """Main entry point for vLLM-Omni inference.
This class provides a high-level interface for running multi-modal
inference with non-autoregressive models.
@@ -50,7 +50,7 @@ class OmniLLM:
**kwargs: Additional arguments passed to the engine
Example:
- >>> llm = OmniLLM(model="Qwen/Qwen2.5-Omni")
+ >>> llm = Omni(model="Qwen/Qwen2.5-Omni")
>>> outputs = llm.generate(prompts="Hello")
"""
```
diff --git a/docs/contributing/README.md b/docs/contributing/README.md
index 43bccdc5b..5f69a3170 100644
--- a/docs/contributing/README.md
+++ b/docs/contributing/README.md
@@ -1,55 +1,49 @@
-# Contributing to vLLM-omni
+# Contributing to vLLM-Omni
-Thank you for your interest in contributing to vLLM-omni! This document provides guidelines and instructions for contributing.
+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:
+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
```
+### Development Environment for vLLM
-Install vLLM from source:
+vLLM-Omni is built on top of the stable release 0.11.0 of vLLM. You may install `vllm==0.11.0` directly or install the library from source if you need to check, modify or debug with source code of vLLM.
-```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 .
-```
+### Development Environment for vLLM-Omni
-Install vLLM-omni from source:
+Install vLLM-Omni from source with development dependencies:
```bash
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm_omni
-uv pip install -e .
+uv pip install -e ".[dev]"
```
!!! 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.
+ 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:
+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.
+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)
+ pre-commit run --show-diff-on-failure --color=always --all-files # runs on all files (short for --all-files)
```
### Documentation
@@ -78,20 +72,14 @@ For additional features and advanced configurations, refer to the:
### Testing
-vLLM-omni uses `pytest` to test the codebase.
+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
+pytest -s -v tests/test_omni_llm.py
```
!!! warning
@@ -106,7 +94,7 @@ If you encounter a bug or have a feature request, please search existing issues
## 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.
+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
@@ -128,9 +116,9 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
- `[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.)
+- `[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.)
+- `[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.
@@ -145,7 +133,7 @@ The PR needs to meet the following code quality standards:
- 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.
+- 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
@@ -153,18 +141,17 @@ Please keep the changes as concise as possible. For major architectural changes
### 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:
+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 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
+- [Design Documents](../design/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!
+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!
diff --git a/docs/contributing/ci/failures.md b/docs/contributing/ci/failures.md
new file mode 100644
index 000000000..d64d98a60
--- /dev/null
+++ b/docs/contributing/ci/failures.md
@@ -0,0 +1,4 @@
+# CI Failures
+
+What should I do when a CI job fails on my PR, but I don't think my PR caused
+the failure?
diff --git a/docs/contributing/model/README.md b/docs/contributing/model/README.md
new file mode 100644
index 000000000..2e7c3a529
--- /dev/null
+++ b/docs/contributing/model/README.md
@@ -0,0 +1,2 @@
+# Adding a New Model
+Stay tuned for guidance on how to add a new model on vLLM-Omni!
diff --git a/docs/contributing/design_documents/api_design_doc.md b/docs/design/api_design_doc.md
similarity index 98%
rename from docs/contributing/design_documents/api_design_doc.md
rename to docs/design/api_design_doc.md
index f67f325fe..b29c44ab5 100644
--- a/docs/contributing/design_documents/api_design_doc.md
+++ b/docs/design/api_design_doc.md
@@ -1,6 +1,6 @@
# 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.
+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
diff --git a/docs/contributing/design_documents/api_design_template.md b/docs/design/api_design_template.md
similarity index 98%
rename from docs/contributing/design_documents/api_design_template.md
rename to docs/design/api_design_template.md
index f9b0478dd..2fe2b233f 100644
--- a/docs/contributing/design_documents/api_design_template.md
+++ b/docs/design/api_design_template.md
@@ -1,6 +1,6 @@
-# API Design Template for vLLM-omni Modules
+# 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.
+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
@@ -44,7 +44,7 @@ This template provides a standardized structure for designing APIs for core, eng
### 1. Module Overview
-**Purpose**: The core module provides fundamental scheduling and caching functionality for vLLM-omni.
+**Purpose**: The core module provides fundamental scheduling and caching functionality for vLLM-Omni.
**Responsibilities**:
- Request scheduling and prioritization
@@ -333,4 +333,4 @@ print(f"Cache hit rate: {status['cache_hit_rate']:.2%}")
6. Include error handling strategies
7. Submit for review before implementation
-This template ensures consistency and completeness across all vLLM-omni modules.
+This template ensures consistency and completeness across all vLLM-Omni modules.
diff --git a/docs/contributing/design_documents/architecture_overview.md b/docs/design/architecture_overview.md
similarity index 99%
rename from docs/contributing/design_documents/architecture_overview.md
rename to docs/design/architecture_overview.md
index 54319bce0..3a219a564 100644
--- a/docs/contributing/design_documents/architecture_overview.md
+++ b/docs/design/architecture_overview.md
@@ -2,11 +2,11 @@
# 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.
+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:
+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.
@@ -57,12 +57,12 @@ According to analysis for current popular open-source models, most of them have
# 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.
+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.
+* 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.
@@ -89,32 +89,6 @@ vLLM-omni will import core modules from existing vLLM framework and introduce se
![][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]:
diff --git a/docs/contributing/design_documents/index.md b/docs/design/index.md
similarity index 51%
rename from docs/contributing/design_documents/index.md
rename to docs/design/index.md
index 42e360ed4..aefefa0c8 100644
--- a/docs/contributing/design_documents/index.md
+++ b/docs/design/index.md
@@ -1,12 +1,13 @@
# Design Documents
-This section contains design documents and architecture specifications for vLLM-omni.
+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)
+- [vLLM-Omni Software Design Document](vllm_omni_design.md)
+- [Multi-Request Streaming](mrs_design.md)
## API Design Documentation
-- [vLLM-omni API Documentation](api_design_doc.md)
+- [vLLM-Omni API Design Documentation](api_design_doc.md)
diff --git a/docs/mrs_design.md b/docs/design/mrs_design.md
similarity index 100%
rename from docs/mrs_design.md
rename to docs/design/mrs_design.md
diff --git a/docs/contributing/design_documents/vllm_omni_design.md b/docs/design/vllm_omni_design.md
similarity index 98%
rename from docs/contributing/design_documents/vllm_omni_design.md
rename to docs/design/vllm_omni_design.md
index 27b62f691..8fe328878 100644
--- a/docs/contributing/design_documents/vllm_omni_design.md
+++ b/docs/design/vllm_omni_design.md
@@ -2,7 +2,7 @@
## 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.
+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
@@ -36,7 +36,7 @@ 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]
+ 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]
@@ -45,7 +45,7 @@ graph TD
G --> H[Multi-stage Processing]
E --> I[Standard vLLM Pipeline]
- subgraph "vLLM-omni Components"
+ subgraph "vLLM-Omni Components"
F
G
H
@@ -68,7 +68,7 @@ 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
+ # 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)
diff --git a/docs/examples/README.md b/docs/examples/README.md
new file mode 100644
index 000000000..4594d2bed
--- /dev/null
+++ b/docs/examples/README.md
@@ -0,0 +1,6 @@
+# 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.
diff --git a/docs/getting_started/installation/.nav.yml b/docs/getting_started/installation/.nav.yml
new file mode 100644
index 000000000..0fcee9a00
--- /dev/null
+++ b/docs/getting_started/installation/.nav.yml
@@ -0,0 +1,4 @@
+nav:
+ - README.md
+ - gpu.md
+ - npu.md
diff --git a/docs/getting_started/installation/README.md b/docs/getting_started/installation/README.md
new file mode 100644
index 000000000..c301e7a12
--- /dev/null
+++ b/docs/getting_started/installation/README.md
@@ -0,0 +1,7 @@
+# Installation
+
+vLLM supports the following hardware platforms:
+
+- [GPU](gpu.md)
+ - [NVIDIA CUDA](gpu.md)
+- [Ascend NPU](npu.md)
diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md
new file mode 100644
index 000000000..4926f93f7
--- /dev/null
+++ b/docs/getting_started/installation/gpu.md
@@ -0,0 +1,40 @@
+# GPU
+
+vLLM-Omni is a Python library that supports the following GPU variants. The library itself mainly contains python implementations for framework and models.
+
+## Requirements
+
+- OS: Linux
+- Python: 3.12
+
+!!! note
+ vLLM-Omni is currently not natively supported on Windows.
+
+=== "NVIDIA CUDA"
+
+ --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:requirements"
+
+## Set up using Python
+
+### Create a new Python environment
+
+--8<-- "docs/getting_started/installation/python_env_setup.inc.md"
+
+### Pre-built wheels
+
+=== "NVIDIA CUDA"
+
+ --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-wheels"
+
+[](){ #build-from-source }
+
+### Build wheel from source
+
+=== "NVIDIA CUDA"
+
+ --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:build-wheel-from-source"
+
+## Set up using Docker
+
+### Pre-built images
+To be released... Stay tuned!
diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md
new file mode 100644
index 000000000..ce394f8a6
--- /dev/null
+++ b/docs/getting_started/installation/gpu/cuda.inc.md
@@ -0,0 +1,75 @@
+# --8<-- [start:requirements]
+
+- GPU: compute capability 7.0 or higher (e.g., V100, T4, RTX20xx, A100, L4, H100, etc.)
+
+# --8<-- [end:requirements]
+# --8<-- [start:set-up-using-python]
+
+vLLM-Omni depends vLLM. So please follow instructions below mainly for vLLM.
+
+!!! note
+ PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See for more details.
+
+In order to be performant, vLLM has to compile many cuda kernels. The compilation unfortunately introduces binary incompatibility with other CUDA versions and PyTorch versions, even for the same PyTorch version with different building configurations.
+
+Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** environment. If either you have a different CUDA version or you want to use an existing PyTorch installation, you need to build vLLM from source. See [build-from-source-vllm](https://docs.vllm.ai/en/stable/getting_started/installation/gpu/#build-wheel-from-source) for more details.
+
+# --8<-- [start:pre-built-wheels]
+
+#### Installation of vLLM
+
+vLLM-Omni is built 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-Omni
+
+```bash
+uv pip install vllm-omni
+```
+
+# --8<-- [end:pre-built-wheels]
+
+# --8<-- [start:build-wheel-from-source]
+
+#### Installation of vLLM
+If you do not need to modify source code of vLLM, you can directly install the stable 0.11.0 release version of the library
+
+```bash
+uv pip install vllm==0.11.0 --torch-backend=auto
+```
+
+#### 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 .
+```
+
+(Optional) Installation of vLLM from source
+If you want to check, modify or debug with source code of vLLM, install the library from source with the following instructions:
+
+```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 .
+```
+🎯 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
-
-- Qwen-omni (Thinker-Talker-Codec structure)
-- HunyunaImage 3.0 (Ongoing)
-- Qwen-Image (Ongoing)
-
-## 📚 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)
diff --git a/docs/mkdocs/hooks/generate_api_readme.py b/docs/mkdocs/hooks/generate_api_readme.py
new file mode 100644
index 000000000..eb7bbfdbe
--- /dev/null
+++ b/docs/mkdocs/hooks/generate_api_readme.py
@@ -0,0 +1,268 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""
+Hook to automatically generate docs/api/README.md from the codebase.
+
+This script scans the vllm_omni module for public classes and functions,
+categorizes them, and generates a summary README file.
+"""
+
+import ast
+import logging
+from pathlib import Path
+
+logger = logging.getLogger("mkdocs")
+
+ROOT_DIR = Path(__file__).parent.parent.parent.parent
+API_README_PATH = ROOT_DIR / "docs" / "api" / "README.md"
+
+# Category mappings: module prefix -> category name and description
+CATEGORIES = {
+ "entrypoints": {
+ "name": "Entry Points",
+ "description": "Main entry points for vLLM-Omni inference and serving.",
+ },
+ "inputs": {
+ "name": "Inputs",
+ "description": "Input data structures for multi-modal inputs.",
+ },
+ "engine": {
+ "name": "Engine",
+ "description": "Engine classes for offline and online inference.",
+ },
+ "core": {
+ "name": "Core",
+ "description": "Core scheduling and caching components.",
+ },
+ "model_executor": {
+ "name": "Model Executor",
+ "description": "Model execution components.",
+ },
+ "config": {
+ "name": "Configuration",
+ "description": "Configuration classes.",
+ },
+ "worker": {
+ "name": "Workers",
+ "description": "Worker classes and model runners for distributed inference.",
+ },
+}
+
+
+class APIVisitor(ast.NodeVisitor):
+ """AST visitor to extract public classes and module-level functions."""
+
+ def __init__(self, module_path: str):
+ self.module_path = module_path
+ self.classes: list[str] = []
+ self.functions: list[str] = []
+ self._class_stack: list[str] = [] # Track nested class definitions
+
+ def visit_ClassDef(self, node: ast.ClassDef) -> None:
+ """Visit class definitions."""
+ if not node.name.startswith("_"):
+ self.classes.append(f"{self.module_path}.{node.name}")
+ # Track that we're entering a class
+ self._class_stack.append(node.name)
+ self.generic_visit(node)
+ # Remove from stack when done visiting
+ self._class_stack.pop()
+
+ def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+ """Visit function definitions - only collect module-level functions."""
+ # Only collect if we're not inside a class (stack is empty)
+ if not self._class_stack and not node.name.startswith("_"):
+ self.functions.append(f"{self.module_path}.{node.name}")
+ self.generic_visit(node)
+
+ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
+ """Visit async function definitions - only collect module-level functions."""
+ # Only collect if we're not inside a class (stack is empty)
+ if not self._class_stack and not node.name.startswith("_"):
+ self.functions.append(f"{self.module_path}.{node.name}")
+ self.generic_visit(node)
+
+
+def parse_file_for_symbols(file_path: Path, module_path: str) -> tuple[list[str], list[str]]:
+ """
+ Parse a Python file and extract public classes and functions.
+
+ Returns:
+ Tuple of (classes, functions)
+ """
+ try:
+ # If this is __init__.py, use parent module path
+ if file_path.name == "__init__.py":
+ # Remove .__init__ from module path
+ if module_path.endswith(".__init__"):
+ module_path = module_path[:-9]
+
+ with open(file_path, encoding="utf-8") as f:
+ content = f.read()
+
+ tree = ast.parse(content, filename=str(file_path))
+ visitor = APIVisitor(module_path)
+ visitor.visit(tree)
+
+ return visitor.classes, visitor.functions
+ except Exception as e:
+ logger.debug(f"Could not parse {file_path}: {e}")
+ return [], []
+
+
+def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
+ """
+ Scan the vllm_omni package and categorize public symbols.
+
+ Returns:
+ Dict mapping category names to lists of symbol full names
+ """
+ categorized: dict[str, list[str]] = {cat["name"]: [] for cat in CATEGORIES.values()}
+
+ try:
+ # Find the package directory
+ package_path = ROOT_DIR / package_name
+ if not package_path.exists():
+ logger.warning(f"Package path not found: {package_path}")
+ return categorized
+
+ # Walk through all Python files
+ for py_file in package_path.rglob("*.py"):
+ # Skip __init__.py and private modules
+ if py_file.name.startswith("_") and py_file.name != "__init__.py":
+ continue
+
+ # Get module path
+ relative_path = py_file.relative_to(ROOT_DIR)
+ module_path = str(relative_path.with_suffix("")).replace("/", ".").replace("\\", ".")
+
+ # Handle __init__.py - use parent module path
+ if py_file.name == "__init__.py":
+ # Remove .__init__ from module path
+ if module_path.endswith(".__init__"):
+ module_path = module_path[:-9]
+
+ # Determine category from module path
+ category = None
+ for prefix, cat_info in CATEGORIES.items():
+ if prefix in module_path:
+ category = cat_info["name"]
+ break
+
+ if not category:
+ continue
+
+ # Parse file for symbols
+ classes, functions = parse_file_for_symbols(py_file, module_path)
+
+ # Filter out internal implementation classes
+ # Skip classes that look like internal components (DiT layers, etc.)
+ internal_patterns = [
+ "Block",
+ "Layer",
+ "Net",
+ "Embedding",
+ "Norm",
+ "Activation",
+ "Solver",
+ "Pooling",
+ "Attention",
+ "MLP",
+ "DecoderLayer",
+ "InputEmbedding",
+ "TimestepEmbedding",
+ "CodecEmbedding",
+ "DownSample",
+ "UpSample",
+ "Res2Net",
+ "SqueezeExcitation",
+ "TimeDelay",
+ "TorchActivation",
+ "SnakeBeta",
+ "SinusPosition",
+ "RungeKutta",
+ "AMPBlock",
+ "AdaLayerNorm",
+ ]
+
+ # Add classes (filter out internal ones)
+ for class_name in classes:
+ class_short_name = class_name.split(".")[-1]
+ # Skip if it matches internal patterns (unless it's a main model class)
+ if any(pattern in class_short_name for pattern in internal_patterns):
+ # But include main model classes
+ if not any(
+ main_class in class_short_name
+ for main_class in [
+ "ForConditionalGeneration",
+ "Model",
+ "Registry",
+ "Worker",
+ "Runner",
+ "Scheduler",
+ "Manager",
+ "Processor",
+ "Config",
+ ]
+ ):
+ continue
+ categorized[category].append(class_name)
+
+ # Add important functions (parse, preprocess, etc.)
+ for func_name in functions:
+ # Include functions that match certain patterns
+ if any(keyword in func_name.lower() for keyword in ["parse", "preprocess"]):
+ categorized[category].append(func_name)
+
+ # Sort symbols within each category
+ for category in categorized:
+ categorized[category].sort()
+
+ except Exception as e:
+ logger.error(f"Error scanning package: {e}", exc_info=True)
+
+ return categorized
+
+
+def generate_readme(categorized: dict[str, list[str]]) -> str:
+ """Generate the API README markdown content."""
+ lines = ["# Summary", ""]
+
+ # Generate sections for each category
+ for prefix, cat_info in CATEGORIES.items():
+ category_name = cat_info["name"]
+ description = cat_info["description"]
+ symbols = categorized.get(category_name, [])
+
+ if not symbols:
+ continue
+
+ lines.append(f"## {category_name}")
+ lines.append("")
+ lines.append(description)
+ lines.append("")
+
+ for symbol in symbols:
+ lines.append(f"- [{symbol}][]")
+
+ lines.append("")
+
+ return "\n".join(lines)
+
+
+def on_startup(command, dirty: bool):
+ """MkDocs hook entry point."""
+ logger.info("Generating API README documentation")
+
+ # Scan the package
+ categorized = scan_package()
+
+ # Generate README content
+ content = generate_readme(categorized)
+
+ # Write to file
+ API_README_PATH.parent.mkdir(parents=True, exist_ok=True)
+ with open(API_README_PATH, "w", encoding="utf-8") as f:
+ f.write(content)
+
+ logger.info(f"API README generated: {API_README_PATH.relative_to(ROOT_DIR)}")
diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py
deleted file mode 100644
index a1f75983f..000000000
--- a/docs/mkdocs/hooks/generate_examples.py
+++ /dev/null
@@ -1,233 +0,0 @@
-# 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.+)$", 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))
diff --git a/docs/mkdocs/hooks/url_schemes.py b/docs/mkdocs/hooks/url_schemes.py
new file mode 100644
index 000000000..8798b11db
--- /dev/null
+++ b/docs/mkdocs/hooks/url_schemes.py
@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+This is basically a port of MyST parser’s external URL resolution mechanism
+(https://myst-parser.readthedocs.io/en/latest/syntax/cross-referencing.html#customising-external-url-resolution)
+to work with MkDocs.
+
+It allows Markdown authors to use GitHub shorthand links like:
+
+ - [Text](gh-issue:123)
+ - <gh-pr:456>
+ - [File](gh-file:path/to/file.py#L10)
+
+These are automatically rewritten into fully qualified GitHub URLs pointing to
+issues, pull requests, files, directories, or projects in the
+`vllm-project/vllm-omni` repository.
+
+The goal is to simplify cross-referencing common GitHub resources
+in project docs.
+"""
+
+import regex as re
+from mkdocs.config.defaults import MkDocsConfig
+from mkdocs.structure.files import Files
+from mkdocs.structure.pages import Page
+
+
+def on_page_markdown(markdown: str, *, page: Page, config: MkDocsConfig, files: Files) -> str:
+ """
+ Custom MkDocs plugin hook to rewrite special GitHub reference links
+ in Markdown.
+
+ This function scans the given Markdown content for specially formatted
+ GitHub shorthand links, such as:
+ - `[Link text](gh-issue:123)`
+ - `<gh-pr:456>`
+
+ And rewrites them into fully-qualified GitHub URLs with GitHub icons:
+ - `[:octicons-mark-github-16: Link text](https://github.com/vllm-project/vllm/issues/123)`
+ - `[:octicons-mark-github-16: Pull Request #456](https://github.com/vllm-project/vllm/pull/456)`
+
+ Supported shorthand types:
+ - `gh-issue`
+ - `gh-pr`
+ - `gh-project`
+ - `gh-dir`
+ - `gh-file`
+
+ Args:
+ markdown (str): The raw Markdown content of the page.
+ page (Page): The MkDocs page object being processed.
+ config (MkDocsConfig): The MkDocs site configuration.
+ files (Files): The collection of files in the MkDocs build.
+
+ Returns:
+ str: The updated Markdown content with GitHub shorthand links replaced.
+ """
+ gh_icon = ":octicons-mark-github-16:"
+ gh_url = "https://github.com"
+ repo_url = f"{gh_url}/vllm-project/vllm-omni"
+ org_url = f"{gh_url}/orgs/vllm-project"
+
+ # Mapping of shorthand types to their corresponding GitHub base URLs
+ urls = {
+ "issue": f"{repo_url}/issues",
+ "pr": f"{repo_url}/pull",
+ "project": f"{org_url}/projects",
+ "dir": f"{repo_url}/tree/main",
+ "file": f"{repo_url}/blob/main",
+ }
+
+ # Default title prefixes for auto links
+ titles = {
+ "issue": "Issue #",
+ "pr": "Pull Request #",
+ "project": "Project #",
+ "dir": "",
+ "file": "",
+ }
+
+ # Regular expression to match GitHub shorthand links
+ scheme = r"gh-(?P<type>.+?):(?P<path>.+?)(#(?P<fragment>.+?))?"
+ inline_link = re.compile(r"\[(?P<title>[^\[]+?)\]\(" + scheme + r"\)")
+ auto_link = re.compile(f"<{scheme}>")
+
+ def replace_inline_link(match: re.Match) -> str:
+ """
+ Replaces a matched inline-style GitHub shorthand link
+ with a full Markdown link.
+
+ Example:
+ [My issue](gh-issue:123) → [:octicons-mark-github-16: My issue](https://github.com/vllm-project/vllm/issues/123)
+ """
+ url = f"{urls[match.group('type')]}/{match.group('path')}"
+ if fragment := match.group("fragment"):
+ url += f"#{fragment}"
+
+ return f"[{gh_icon} {match.group('title')}]({url})"
+
+ def replace_auto_link(match: re.Match) -> str:
+ """
+ Replaces a matched autolink-style GitHub shorthand
+ with a full Markdown link.
+
+ Example:
+ <gh-pr:456> → [:octicons-mark-github-16: Pull Request #456](https://github.com/vllm-project/vllm/pull/456)
+ """
+ type = match.group("type")
+ path = match.group("path")
+ title = f"{titles[type]}{path}"
+ url = f"{urls[type]}/{path}"
+ if fragment := match.group("fragment"):
+ url += f"#{fragment}"
+
+ return f"[{gh_icon} {title}]({url})"
+
+ # Replace both inline and autolinks
+ markdown = inline_link.sub(replace_inline_link, markdown)
+ markdown = auto_link.sub(replace_auto_link, markdown)
+
+ return markdown
diff --git a/docs/mkdocs/javascript/edit_and_feedback.js b/docs/mkdocs/javascript/edit_and_feedback.js
new file mode 100644
index 000000000..0acbcfd96
--- /dev/null
+++ b/docs/mkdocs/javascript/edit_and_feedback.js
@@ -0,0 +1,47 @@
+/**
+ * edit_and_feedback.js
+ *
+ * Enhances MkDocs Material docs pages by:
+ *
+ * 1. Adding a "Question? Give us feedback" link
+ * below the "Edit" button.
+ *
+ * - The link opens a GitHub issue with a template,
+ * auto-filled with the current page URL and path.
+ *
+ * 2. Ensuring the edit button opens in a new tab
+ * with target="_blank" and rel="noopener".
+ */
+document.addEventListener("DOMContentLoaded", function () {
+ const url = window.location.href;
+ const page = document.body.dataset.mdUrl || location.pathname;
+
+ const feedbackLink = document.createElement("a");
+ feedbackLink.href = `https://github.com/vllm-project/vllm-omni/issues/new?template=100-documentation.yml&title=${encodeURIComponent(
+ `[Docs] Feedback for \`${page}\``
+ )}&body=${encodeURIComponent(`📄 **Reference:**\n${url}\n\n📝 **Feedback:**\n_Your response_`)}`;
+ feedbackLink.target = "_blank";
+ feedbackLink.rel = "noopener";
+ feedbackLink.title = "Provide feedback";
+ feedbackLink.className = "md-content__button";
+ feedbackLink.innerHTML = `
+ <svg
+ xmlns="http://www.w3.org/2000/svg"
+ height="24px"
+ viewBox="0 -960 960 960"
+ width="24px"
+ fill="currentColor"
+ >
+ <path d="M280-280h280v-80H280v80Zm0-160h400v-80H280v80Zm0-160h400v-80H280v80Zm-80 480q-33 0-56.5-23.5T120-200v-560q0-33 23.5-56.5T200-840h560q33 0 56.5 23.5T840-760v560q0 33-23.5 56.5T760-120H200Zm0-80h560v-560H200v560Zm0-560v560-560Z"/>
+ </svg>
+ `;
+
+ const editButton = document.querySelector('.md-content__button[href*="edit"]');
+
+ if (editButton && editButton.parentNode) {
+ editButton.insertAdjacentElement("beforebegin", feedbackLink);
+
+ editButton.setAttribute("target", "_blank");
+ editButton.setAttribute("rel", "noopener");
+ }
+ });
diff --git a/docs/mkdocs/javascript/mathjax.js b/docs/mkdocs/javascript/mathjax.js
new file mode 100644
index 000000000..eb89ace06
--- /dev/null
+++ b/docs/mkdocs/javascript/mathjax.js
@@ -0,0 +1,20 @@
+// Enables MathJax rendering
+window.MathJax = {
+ tex: {
+ inlineMath: [["\\(", "\\)"]],
+ displayMath: [["\\[", "\\]"]],
+ processEscapes: true,
+ processEnvironments: true
+ },
+ options: {
+ ignoreHtmlClass: ".*|",
+ processHtmlClass: "arithmatex"
+ }
+ };
+
+ document$.subscribe(() => {
+ MathJax.startup.output.clearCache()
+ MathJax.typesetClear()
+ MathJax.texReset()
+ MathJax.typesetPromise()
+ })
diff --git a/docs/mkdocs/javascript/slack_and_forum.js b/docs/mkdocs/javascript/slack_and_forum.js
new file mode 100644
index 000000000..320e44910
--- /dev/null
+++ b/docs/mkdocs/javascript/slack_and_forum.js
@@ -0,0 +1,56 @@
+/**
+ * slack_and_forum.js
+ *
+ * Adds a custom Slack and Forum button to the MkDocs Material header.
+ *
+ */
+
+window.addEventListener('DOMContentLoaded', () => {
+ const headerInner = document.querySelector('.md-header__inner');
+
+ if (headerInner) {
+ const slackButton = document.createElement('button');
+ slackButton.className = 'slack-button';
+ slackButton.title = 'Join us on Slack';
+ slackButton.style.border = 'none';
+ slackButton.style.background = 'transparent';
+ slackButton.style.cursor = 'pointer';
+
+ slackButton.innerHTML = `
+ <img src="https://a.slack-edge.com/80588/marketing/img/icons/icon_slack_hash_colored.png"
+ style="height: 1.1rem;"
+ alt="Slack">
+ `;
+
+ slackButton.addEventListener('click', () => {
+ window.open('https://slack.vllm.ai', '_blank', 'noopener');
+ });
+
+ const forumButton = document.createElement('button');
+ forumButton.className = 'forum-button';
+ forumButton.title = 'Join the Forum';
+ forumButton.style.border = 'none';
+ forumButton.style.background = 'transparent';
+ forumButton.style.cursor = 'pointer';
+
+ forumButton.innerHTML = `
+ <svg
+ xmlns="http://www.w3.org/2000/svg"
+ viewBox="0 -960 960 960"
+ fill="currentColor"
+ >
+ <path d="M817.85-198.15 698.46-317.54H320q-24.48 0-41.47-16.99T261.54-376v-11.69h424.61q25.39 0 43.47-18.08 18.07-18.08 18.07-43.46v-268.92h11.69q24.48 0 41.47 16.99 17 16.99 17 41.47v461.54ZM179.08-434.69l66.84-66.85h363.31q10.77 0 17.69-6.92 6.93-6.92 6.93-17.69v-246.77q0-10.77-6.93-17.7-6.92-6.92-17.69-6.92H203.69q-10.77 0-17.69 6.92-6.92 6.93-6.92 17.7v338.23Zm-36.93 89.46v-427.69q0-25.39 18.08-43.46 18.08-18.08 43.46-18.08h405.54q25.39 0 43.46 18.08 18.08 18.07 18.08 43.46v246.77q0 25.38-18.08 43.46-18.07 18.07-43.46 18.07H261.54L142.15-345.23Zm36.93-180.92V-797.54v271.39Z"/>
+ </svg>
+ `;
+
+ forumButton.addEventListener('click', () => {
+ window.open('https://discuss.vllm.ai/', '_blank', 'noopener');
+ });
+
+ const githubSource = document.querySelector('.md-header__source');
+ if (githubSource) {
+ githubSource.parentNode.insertBefore(slackButton, githubSource.nextSibling);
+ githubSource.parentNode.insertBefore(forumButton, slackButton.nextSibling);
+ }
+ }
+ });
diff --git a/docs/mkdocs/overrides/main.html b/docs/mkdocs/overrides/main.html
new file mode 100644
index 000000000..94d9808cc
--- /dev/null
+++ b/docs/mkdocs/overrides/main.html
@@ -0,0 +1 @@
+{% extends "base.html" %}
diff --git a/docs/mkdocs/overrides/partials/toc-item.html b/docs/mkdocs/overrides/partials/toc-item.html
new file mode 100644
index 000000000..656c21589
--- /dev/null
+++ b/docs/mkdocs/overrides/partials/toc-item.html
@@ -0,0 +1,21 @@
+<!-- Enables the use of toc_depth in document frontmatter https://github.com/squidfunk/mkdocs-material/issues/4827#issuecomment-1869812019 -->
+<li class="md-nav__item">
+ <a href="{{ toc_item.url }}" class="md-nav__link">
+ <span class="md-ellipsis">
+ {{ toc_item.title }}
+ </span>
+ </a>
+
+ <!-- Table of contents list -->
+ {% if toc_item.children %}
+ <nav class="md-nav" aria-label="{{ toc_item.title | striptags }}">
+ <ul class="md-nav__list">
+ {% for toc_item in toc_item.children %}
+ {% if not page.meta.toc_depth or toc_item.level <= page.meta.toc_depth %}
+ {% include "partials/toc-item.html" %}
+ {% endif %}
+ {% endfor %}
+ </ul>
+ </nav>
+ {% endif %}
+ </li>
diff --git a/docs/mkdocs/stylesheets/extra.css b/docs/mkdocs/stylesheets/extra.css
index efc636b53..5f6ec03d6 100644
--- a/docs/mkdocs/stylesheets/extra.css
+++ b/docs/mkdocs/stylesheets/extra.css
@@ -1,55 +1,218 @@
-/* Custom CSS for vLLM-omni documentation - based on vLLM's extra.css */
-
-/* Improve code block readability */
-.highlight {
- border-radius: 4px;
+/* Warning for latest docs */
+.md-banner {
+ background-color: var(--md-warning-bg-color);
+ color: var(--md-warning-fg-color);
}
-/* Style API documentation */
-.md-typeset .doc-contents {
- margin-bottom: 1.5em;
+/* https://christianoliff.com/blog/styling-external-links-with-an-icon-in-css/ */
+a:not(:has(svg)):not(.md-icon):not(.autorefs-external) {
+ align-items: center;
+
+ &[href^="//"]::after,
+ &[href^="http://"]::after,
+ &[href^="https://"]::after {
+ content: "";
+ width: 12px;
+ height: 12px;
+ margin-left: 4px;
+ background-image: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='16' height='16' stroke='gray' viewBox='0 0 16 16'%3E%3Cpath fill-rule='evenodd' d='M8.636 3.5a.5.5 0 0 0-.5-.5H1.5A1.5 1.5 0 0 0 0 4.5v10A1.5 1.5 0 0 0 1.5 16h10a1.5 1.5 0 0 0 1.5-1.5V7.864a.5.5 0 0 0-1 0V14.5a.5.5 0 0 1-.5.5h-10a.5.5 0 0 1-.5-.5v-10a.5.5 0 0 1 .5-.5h6.636a.5.5 0 0 0 .5-.5z'/%3E%3Cpath fill-rule='evenodd' d='M16 .5a.5.5 0 0 0-.5-.5h-5a.5.5 0 0 0 0 1h3.793L6.146 9.146a.5.5 0 1 0 .708.708L15 1.707V5.5a.5.5 0 0 0 1 0v-5z'/%3E%3C/svg%3E");
+ background-position: center;
+ background-repeat: no-repeat;
+ background-size: contain;
+ display: inline-block;
+ }
}
-/* Improve table readability */
-.md-typeset table {
- display: table;
- width: 100%;
+a[href*="localhost"]::after,
+a[href*="127.0.0.1"]::after,
+a[href*="org.readthedocs.build"]::after,
+a[href*="docs.vllm.ai"]::after {
+ display: none !important;
}
-/* Light mode: darker section titles (from vLLM) */
+/* Light mode: darker section titles */
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) */
+/* Dark mode: lighter gray section titles */
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;
+/* Custom admonitions */
+:root {
+ --md-admonition-icon--announcement: url('data:image/svg+xml;charset=utf-8,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" width="16" height="16"><path d="M3.25 9a.75.75 0 0 1 .75.75c0 2.142.456 3.828.733 4.653a.122.122 0 0 0 .05.064.212.212 0 0 0 .117.033h1.31c.085 0 .18-.042.258-.152a.45.45 0 0 0 .075-.366A16.743 16.743 0 0 1 6 9.75a.75.75 0 0 1 1.5 0c0 1.588.25 2.926.494 3.85.293 1.113-.504 2.4-1.783 2.4H4.9c-.686 0-1.35-.41-1.589-1.12A16.4 16.4 0 0 1 2.5 9.75.75.75 0 0 1 3.25 9Z"></path><path d="M0 6a4 4 0 0 1 4-4h2.75a.75.75 0 0 1 .75.75v6.5a.75.75 0 0 1-.75.75H4a4 4 0 0 1-4-4Zm4-2.5a2.5 2.5 0 1 0 0 5h2v-5Z"></path><path d="M15.59.082A.75.75 0 0 1 16 .75v10.5a.75.75 0 0 1-1.189.608l-.002-.001h.001l-.014-.01a5.775 5.775 0 0 0-.422-.25 10.63 10.63 0 0 0-1.469-.64C11.576 10.484 9.536 10 6.75 10a.75.75 0 0 1 0-1.5c2.964 0 5.174.516 6.658 1.043.423.151.787.302 1.092.443V2.014c-.305.14-.669.292-1.092.443C11.924 2.984 9.713 3.5 6.75 3.5a.75.75 0 0 1 0-1.5c2.786 0 4.826-.484 6.155-.957.665-.236 1.154-.47 1.47-.64.144-.077.284-.161.421-.25l.014-.01a.75.75 0 0 1 .78-.061Z"></path></svg>');
+ --md-admonition-icon--important: url('data:image/svg+xml;charset=utf-8,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" width="16" height="16"><path d="M4.47.22A.749.749 0 0 1 5 0h6c.199 0 .389.079.53.22l4.25 4.25c.141.14.22.331.22.53v6a.749.749 0 0 1-.22.53l-4.25 4.25A.749.749 0 0 1 11 16H5a.749.749 0 0 1-.53-.22L.22 11.53A.749.749 0 0 1 0 11V5c0-.199.079-.389.22-.53Zm.84 1.28L1.5 5.31v5.38l3.81 3.81h5.38l3.81-3.81V5.31L10.69 1.5ZM8 4a.75.75 0 0 1 .75.75v3.5a.75.75 0 0 1-1.5 0v-3.5A.75.75 0 0 1 8 4Zm0 8a1 1 0 1 1 0-2 1 1 0 0 1 0 2Z"></path></svg>');
+ --md-admonition-icon--code: url('data:image/svg+xml;charset=utf-8,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16"><path d="m11.28 3.22 4.25 4.25a.75.75 0 0 1 0 1.06l-4.25 4.25a.749.749 0 0 1-1.275-.326.75.75 0 0 1 .215-.734L13.94 8l-3.72-3.72a.749.749 0 0 1 .326-1.275.75.75 0 0 1 .734.215m-6.56 0a.75.75 0 0 1 1.042.018.75.75 0 0 1 .018 1.042L2.06 8l3.72 3.72a.749.749 0 0 1-.326 1.275.75.75 0 0 1-.734-.215L.47 8.53a.75.75 0 0 1 0-1.06Z"/></svg>');
+ --md-admonition-icon--console: url('data:image/svg+xml;charset=utf-8,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16"><path d="M0 2.75C0 1.784.784 1 1.75 1h12.5c.966 0 1.75.784 1.75 1.75v10.5A1.75 1.75 0 0 1 14.25 15H1.75A1.75 1.75 0 0 1 0 13.25Zm1.75-.25a.25.25 0 0 0-.25.25v10.5c0 .138.112.25.25.25h12.5a.25.25 0 0 0 .25-.25V2.75a.25.25 0 0 0-.25-.25ZM7.25 8a.75.75 0 0 1-.22.53l-2.25 2.25a.749.749 0 0 1-1.275-.326.75.75 0 0 1 .215-.734L5.44 8 3.72 6.28a.749.749 0 0 1 .326-1.275.75.75 0 0 1 .734.215l2.25 2.25c.141.14.22.331.22.53m1.5 1.5h3a.75.75 0 0 1 0 1.5h-3a.75.75 0 0 1 0-1.5"/></svg>');
}
-/* Expand active items */
-.md-nav__item--nested.md-nav__item--active > .md-nav {
- display: block;
+.md-typeset .admonition.announcement,
+.md-typeset details.announcement {
+ border-color: rgb(255, 110, 66);
+}
+.md-typeset .admonition.important,
+.md-typeset details.important {
+ border-color: rgb(239, 85, 82);
+}
+.md-typeset .admonition.code,
+.md-typeset details.code {
+ border-color: #64dd17
+}
+.md-typeset .admonition.console,
+.md-typeset details.console {
+ border-color: #64dd17
}
-/* 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;
+.md-typeset .announcement > .admonition-title,
+.md-typeset .announcement > summary {
+ background-color: rgb(255, 110, 66, 0.1);
+}
+.md-typeset .important > .admonition-title,
+.md-typeset .important > summary {
+ background-color: rgb(239, 85, 82, 0.1);
+}
+.md-typeset .code > .admonition-title,
+.md-typeset .code > summary {
+ background-color: #64dd171a;
+}
+.md-typeset .console > .admonition-title,
+.md-typeset .console > summary {
+ background-color: #64dd171a;
}
-/* 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;
+.md-typeset .announcement > .admonition-title::before,
+.md-typeset .announcement > summary::before {
+ background-color: rgb(239, 85, 82);
+ -webkit-mask-image: var(--md-admonition-icon--announcement);
+ mask-image: var(--md-admonition-icon--announcement);
+}
+.md-typeset .important > .admonition-title::before,
+.md-typeset .important > summary::before {
+ background-color: rgb(239, 85, 82);
+ -webkit-mask-image: var(--md-admonition-icon--important);
+ mask-image: var(--md-admonition-icon--important);
+}
+.md-typeset .code > .admonition-title::before,
+.md-typeset .code > summary::before {
+ background-color: #64dd17;
+ -webkit-mask-image: var(--md-admonition-icon--code);
+ mask-image: var(--md-admonition-icon--code);
+}
+.md-typeset .console > .admonition-title::before,
+.md-typeset .console > summary::before {
+ background-color: #64dd17;
+ -webkit-mask-image: var(--md-admonition-icon--console);
+ mask-image: var(--md-admonition-icon--console);
}
-/* 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;
+/* Make label fully visible on hover */
+.md-content__button[href*="edit"]:hover::after {
+ opacity: 1;
+}
+
+/* Hide edit button on generated docs/examples pages */
+@media (min-width: 960px) {
+ .md-content__button[href*="docs/examples/"] {
+ display: none !important;
+ }
+}
+
+.md-content__button-wrapper {
+ position: absolute;
+ top: 0.6rem;
+ right: 0.8rem;
+ display: flex;
+ flex-direction: row;
+ align-items: center;
+ gap: 0.4rem;
+ z-index: 1;
+}
+
+.md-content__button-wrapper a {
+ display: inline-flex;
+ align-items: center;
+ justify-content: center;
+ height: 24px;
+ width: 24px;
+ color: var(--md-default-fg-color);
+ text-decoration: none;
+}
+
+.md-content__button-wrapper a:hover {
+ color: var(--md-accent-fg-color);
+}
+
+/* Slack and Forum css */
+.slack-button,
+.forum-button {
+ display: inline-flex;
+ align-items: center;
+ justify-content: center;
+ margin-left: 0.4rem;
+ height: 24px;
+}
+
+.slack-button img {
+ height: 18px;
+ filter: none !important;
+}
+
+.slack-button:hover,
+.forum-button:hover {
+ opacity: 0.7;
+}
+
+.forum-button svg {
+ height: 28px;
+ opacity: 0.9;
+ transform: translateY(2px);
+}
+
+/* For logo css */
+[data-md-color-scheme="default"] .logo-dark {
+ display: none;
+}
+
+[data-md-color-scheme="slate"] .logo-light {
+ display: none;
+}
+
+/* Outline for content tabs */
+.md-typeset .tabbed-set {
+ border: 0.075rem solid var(--md-default-fg-color);
+ border-radius: 0.2rem;
+}
+
+.md-typeset .tabbed-content {
+ padding: 0 0.6em;
+}
+
+/* Hide link icon in header logo */
+.md-header__button.md-logo :is(img, svg) {
+ pointer-events: none;
+}
+
+.md-header__button.md-logo::after {
+ display: none !important;
+}
+
+/* Hide link icons in content tabs (tabbed content) */
+.md-typeset .tabbed-set > label::after,
+.md-typeset .tabbed-set > input:checked + label::after {
+ display: none !important;
+}
+
+/* Hide link icons in navigation tabs */
+.md-nav__link[href]::after,
+.md-nav__item--nested > .md-nav__link::after {
+ display: none !important;
+}
+
+/* Hide link icons in top navigation tabs */
+.md-tabs__link::after {
+ display: none !important;
}
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
new file mode 100644
index 000000000..a12f2da42
--- /dev/null
+++ b/docs/models/supported_models.md
@@ -0,0 +1,35 @@
+# Supported Models
+
+vLLM-Omni supports unified multimodal comprehension and generation models across various tasks.
+
+## Model Implementation
+
+If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm-omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
+
+## List of Supported Models for Nvidia GPU
+
+<style>
+th {
+ white-space: nowrap;
+ min-width: 0 !important;
+}
+</style>
+
+| Architecture | Models | Example HF Models |
+|--------------|--------|-------------------|
+| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
+| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
+
+
+## List of Supported Models for Ascend NPU
+
+<style>
+th {
+ white-space: nowrap;
+ min-width: 0 !important;
+}
+</style>
+
+| Architecture | Models | Example HF Models |
+|--------------|--------|-------------------|
+| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
diff --git a/docs/source/logos/vllm-logo-only-light.ico b/docs/source/logos/vllm-logo-only-light.ico
new file mode 100644
index 000000000..27528ceeb
Binary files /dev/null and b/docs/source/logos/vllm-logo-only-light.ico differ
diff --git a/docs/usage/faq.md b/docs/usage/faq.md
new file mode 100644
index 000000000..88a62bda9
--- /dev/null
+++ b/docs/usage/faq.md
@@ -0,0 +1,13 @@
+# Frequently Asked Questions
+
+> Q: How many chips do I need to infer a model in vLLM-Omni?
+
+A: Now, we support natively disaggregated deployment for different model stages within a model. There is a restriction that one chip can only has one AutoRegressive model stage. This is because the unified KV cache management of vLLM. Stages of other types can coexist with in a chip. The restriction will be resolved in later version.
+
+> Q: When trying to run examples, I encounter error about backend of librosa or soundfile. How to solve it?
+
+A: If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
diff --git a/docs/user_guide/examples/index.md b/docs/user_guide/examples/index.md
deleted file mode 100644
index d6e084006..000000000
--- a/docs/user_guide/examples/index.md
+++ /dev/null
@@ -1,8 +0,0 @@
-# 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.
diff --git a/docs/user_guide/examples/offline_inference/qwen2_5_omni.md b/docs/user_guide/examples/offline_inference/qwen2_5_omni.md
index 576338003..c11360bcf 100644
--- a/docs/user_guide/examples/offline_inference/qwen2_5_omni.md
+++ b/docs/user_guide/examples/offline_inference/qwen2_5_omni.md
@@ -1,64 +1,60 @@
-# Offline Example of vLLM-omni for Qwen2.5-omni
+# Offline Inference Example of vLLM-Omni for Qwen2.5-Omni
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples\offline_inference\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)
+Please refer to [installation](../../../getting_started/installation/README.md).
+
+## Run Offline inference with Qwen2.5-Omni.
+First, navigate to the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Inside the directory, `end2end.py` is a comprehensive demo suite for initializing a model instance of `Qwen/Qwen2.5-Omni-7B` and use it for various offline inference tasks.
+??? abstract "end2end.py"
+ ``````py
+ --8<-- "examples/offline_inference/qwen2_5_omni/end2end.py"
+ ``````
+
+Below we also provide simple bash scripts to execute this file.
+### Single Prompt
+```bash
+bash run_single_prompt.sh
+```
+??? abstract "run_single_prompt.sh"
+ ``````sh
+ --8<-- "examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
+ ``````
-## 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:
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit) and extract the prompts with `extract_prompts.py`
+??? abstract "extract_prompts.py"
+ ``````py
+ --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
+ ``````
```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
+ --topk 10
```
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"
+ --8<-- "examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
``````
+
+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
diff --git a/docs/user_guide/examples/offline_inference/qwen_2_5_omni.md b/docs/user_guide/examples/offline_inference/qwen_2_5_omni.md
deleted file mode 100644
index 0f9aa971f..000000000
--- a/docs/user_guide/examples/offline_inference/qwen_2_5_omni.md
+++ /dev/null
@@ -1,97 +0,0 @@
-# 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"
- ``````
diff --git a/docs/user_guide/examples/online_serving/README.md b/docs/user_guide/examples/online_serving/README.md
deleted file mode 100644
index 7b3341255..000000000
--- a/docs/user_guide/examples/online_serving/README.md
+++ /dev/null
@@ -1,35 +0,0 @@
-# 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
-```
diff --git a/docs/user_guide/examples/online_serving/gradio_demo.md b/docs/user_guide/examples/online_serving/gradio_demo.md
new file mode 100644
index 000000000..38278d9cf
--- /dev/null
+++ b/docs/user_guide/examples/online_serving/gradio_demo.md
@@ -0,0 +1,7 @@
+# Gradio Demo
+
+Source <https://github.com/vllm-project/vllm-omni/blob/main/examples/online_serving/gradio_demo.py>.
+
+``````py
+--8<-- "examples/online_serving/gradio_demo.py"
+``````
diff --git a/docs/user_guide/examples/online_serving/index.md b/docs/user_guide/examples/online_serving/index.md
deleted file mode 100644
index 6722ed568..000000000
--- a/docs/user_guide/examples/online_serving/index.md
+++ /dev/null
@@ -1,5 +0,0 @@
-# Online Serving
-
-## Available Examples
-
-Online serving is under active development. Coming soon.
diff --git a/docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md b/docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md
index 066b27624..ca3fa8306 100644
--- a/docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md
+++ b/docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md
@@ -1,7 +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>.
+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"
+--8<-- "examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py"
``````
diff --git a/docs/user_guide/examples/online_serving/qwen2_5_omni.md b/docs/user_guide/examples/online_serving/qwen2_5_omni.md
new file mode 100644
index 000000000..96c283d16
--- /dev/null
+++ b/docs/user_guide/examples/online_serving/qwen2_5_omni.md
@@ -0,0 +1,80 @@
+# 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 [installation](../../../getting_started/installation/README.md).
+
+## Deploy Qwen/Qwen2.5-Omni-7B
+
+First launch the OpenAI-compatible inference 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
+```
+## Query the model
+Navigate to the example folder
+```bash
+cd examples/online_serving
+```
+Query the model server with OpenAI Python API client:
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
+```
+??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
+ ``````py
+ --8<-- "examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py"
+ ``````
+You can also query the model with `curl` command:
+```bash
+bash run_curl_multimodal_generation.sh mixed_modalities
+```
+??? abstract "run_curl_multimodal_generation.sh"
+ ``````py
+ --8<-- "examples/online_serving/run_curl_multimodal_generation.sh"
+ ``````
+
+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
+
+## Run Local Web UI Demo
+
+You can also deploy a Gradio Web UI that allows users to interact with the model through a web browser. Below is an example on how to do so with `Qwen/Qwen2.5-Omni-7B`.
+
+### Running Gradio Demo
+
+Install gradio with `uv pip install "gradio>=5.49.1,<6.0.0"`, then you can launch the web service built on AsyncOmni by
+
+```bash
+python gradio_demo.py --model Qwen/Qwen2.5-Omni-7B --port 7861
+```
+Now you can interact with model via the web UI at `http://localhost:7861/` on your local browser.
+
+
+### Options
+
+You can customize its basic launch parameters:
+
+```bash
+python gradio_demo.py \
+ --model Qwen/Qwen2.5-Omni-7B \
+ --ip 127.0.0.1 \
+ --port 7861 \
+ --stage-configs-path /path/to/stage_configs.yaml
+```
+
+- `--model`: Local model checkpoint to load (default `Qwen/Qwen2.5-Omni-7B`).
+- `--ip`: Host/IP for the Gradio server (default `127.0.0.1`).
+- `--port`: Port for the Gradio server (default `7861`).
+- `--stage-configs-path`: Optional path to custom stage configs YAML.
+- `--share`: Set to expose a temporary public link via Gradio.
diff --git a/docs/user_guide/getting_started/quickstart.md b/docs/user_guide/getting_started/quickstart.md
deleted file mode 100644
index 8207b23e8..000000000
--- a/docs/user_guide/getting_started/quickstart.md
+++ /dev/null
@@ -1,55 +0,0 @@
-# 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)
diff --git a/examples/offline_inference/qwen2_5_omni/README.md b/examples/offline_inference/qwen2_5_omni/README.md
index 142039789..afdaee13e 100644
--- a/examples/offline_inference/qwen2_5_omni/README.md
+++ b/examples/offline_inference/qwen2_5_omni/README.md
@@ -1,4 +1,4 @@
-# Offline Example of vLLM-omni for Qwen2.5-omni
+# Offline Example of vLLM-Omni for Qwen2.5-omni
## 🛠️ Installation
diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py
index be4e2c1a1..c5be18a94 100644
--- a/examples/offline_inference/qwen2_5_omni/end2end.py
+++ b/examples/offline_inference/qwen2_5_omni/end2end.py
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
-This example shows how to use vLLM-omni for running offline inference
+This example shows how to use vLLM-Omni for running offline inference
with the correct prompt format on Qwen2.5-Omni
"""
@@ -16,7 +16,7 @@ from vllm.multimodal.image import convert_image_mode
from vllm.sampling_params import SamplingParams
from vllm.utils import FlexibleArgumentParser
-from vllm_omni import OmniLLM
+from vllm_omni.entrypoints.omni_llm import OmniLLM
SEED = 42
diff --git a/examples/offline_inference/qwen3_omni/README.md b/examples/offline_inference/qwen3_omni/README.md
new file mode 100644
index 000000000..766453667
--- /dev/null
+++ b/examples/offline_inference/qwen3_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 (Qwen3-omni)
+### Multiple Prompts
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-omni README.md](../qwen2_5_omni/README.md)
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen3_omni
+```
+Then run the command below.
+```bash
+bash run_multiple_prompts.sh
+```
+### Single Prompts
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen3_omni
+```
+Then run the command below.
+```bash
+bash run_single_prompt.sh
+```
+
+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py
new file mode 100644
index 000000000..c1fa21a62
--- /dev/null
+++ b/examples/offline_inference/qwen3_omni/end2end.py
@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+This example shows how to use vLLM for running offline inference
+with the correct prompt format on Qwen3-Omni (thinker only).
+"""
+
+import os
+from typing import NamedTuple
+
+import soundfile as sf
+from vllm import SamplingParams
+from vllm.assets.audio import AudioAsset
+from vllm.assets.image import ImageAsset
+from vllm.assets.video import VideoAsset
+from vllm.multimodal.image import convert_image_mode
+from vllm.utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+SEED = 42
+
+
+class QueryResult(NamedTuple):
+ inputs: dict
+ limit_mm_per_prompt: dict[str, int]
+
+
+# NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on
+# lower-end GPUs.
+# Unless specified, these settings have been tested to work on a single L4.
+
+default_system = (
+ "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 get_text_query(question: str = None) -> QueryResult:
+ if question is None:
+ question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+ prompt = (
+ f"<|im_start|>system\n{default_system}<|im_end|>\n"
+ "<|im_start|>user\n"
+ f"{question}<|im_end|>\n"
+ f"<|im_start|>assistant\n"
+ )
+ return QueryResult(
+ inputs={
+ "prompt": prompt,
+ },
+ limit_mm_per_prompt={},
+ )
+
+
+def get_video_query(question: str = None) -> QueryResult:
+ if question is None:
+ question = "Why is this video funny?"
+ prompt = (
+ f"<|im_start|>system\n{default_system}<|im_end|>\n"
+ "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
+ f"{question}<|im_end|>\n"
+ f"<|im_start|>assistant\n"
+ )
+
+ return QueryResult(
+ inputs={
+ "prompt": prompt,
+ "multi_modal_data": {
+ "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
+ },
+ },
+ limit_mm_per_prompt={"video": 1},
+ )
+
+
+def get_image_query(question: str = None) -> QueryResult:
+ if question is None:
+ question = "What is the content of this image?"
+ prompt = (
+ f"<|im_start|>system\n{default_system}<|im_end|>\n"
+ "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+ f"{question}<|im_end|>\n"
+ f"<|im_start|>assistant\n"
+ )
+ return QueryResult(
+ inputs={
+ "prompt": prompt,
+ "multi_modal_data": {
+ "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
+ },
+ },
+ limit_mm_per_prompt={"image": 1},
+ )
+
+
+def get_audio_query(question: str = None) -> QueryResult:
+ if question is None:
+ question = "What is the content of this audio?"
+ prompt = (
+ f"<|im_start|>system\n{default_system}<|im_end|>\n"
+ "<|im_start|>user\n<|audio_start|><|audio_pad|><|audio_end|>"
+ f"{question}<|im_end|>\n"
+ f"<|im_start|>assistant\n"
+ )
+ return QueryResult(
+ inputs={
+ "prompt": prompt,
+ "multi_modal_data": {
+ "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
+ },
+ },
+ limit_mm_per_prompt={"audio": 1},
+ )
+
+
+query_map = {
+ "text": get_text_query,
+ "use_audio": get_audio_query,
+ "use_image": get_image_query,
+ "use_video": get_video_query,
+}
+
+
+def main(args):
+ model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+ query_result = query_map[args.query_type]()
+
+ omni_llm = OmniLLM(
+ model=model_name,
+ )
+
+ thinker_sampling_params = SamplingParams(
+ temperature=0.4,
+ top_p=0.9,
+ top_k=-1,
+ max_tokens=1200,
+ repetition_penalty=1.05,
+ logit_bias={},
+ seed=SEED,
+ )
+
+ talker_sampling_params = SamplingParams(
+ temperature=0.9,
+ top_k=50,
+ max_tokens=4096,
+ seed=SEED,
+ detokenize=False,
+ repetition_penalty=1.05,
+ stop_token_ids=[2150], # TALKER_CODEC_EOS_TOKEN_ID
+ )
+
+ # Sampling parameters for Code2Wav stage (audio generation)
+ code2wav_sampling_params = SamplingParams(
+ temperature=0.0,
+ top_p=1.0,
+ top_k=-1,
+ max_tokens=4096 * 16,
+ seed=SEED,
+ detokenize=True,
+ repetition_penalty=1.1,
+ )
+
+ sampling_params_list = [
+ thinker_sampling_params,
+ talker_sampling_params, # code predictor is integrated into talker for Qwen3 Omni
+ code2wav_sampling_params,
+ ]
+
+ if args.txt_prompts is None:
+ prompts = [query_result.inputs for _ in range(args.num_prompts)]
+ else:
+ assert args.query_type == "text", "txt-prompts is only supported for text query type"
+ with open(args.txt_prompts, encoding="utf-8") as f:
+ lines = [ln.strip() for ln in f.readlines()]
+ prompts = [get_text_query(ln).inputs for ln in lines if ln != ""]
+ print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
+
+ omni_outputs = omni_llm.generate(prompts, 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 = prompts[request_id]["prompt"]
+ 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")
+
+ # Convert to numpy array and ensure correct format
+ audio_numpy = audio_tensor.float().detach().cpu().numpy()
+
+ # Ensure audio is 1D (flatten if needed)
+ if audio_numpy.ndim > 1:
+ audio_numpy = audio_numpy.flatten()
+
+ # Save audio file with explicit WAV format
+ sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+ print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+
+
+def parse_args():
+ parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+ parser.add_argument(
+ "--query-type",
+ "-q",
+ type=str,
+ default="mixed_modalities",
+ choices=query_map.keys(),
+ help="Query type.",
+ )
+ parser.add_argument(
+ "--enable-stats",
+ action="store_true",
+ default=False,
+ help="Enable writing detailed statistics (default: disabled)",
+ )
+ 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(
+ "--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(
+ "--output-wav",
+ default="output_audio",
+ help="[Deprecated] Output wav directory (use --output-dir).",
+ )
+ parser.add_argument(
+ "--num-prompts",
+ type=int,
+ default=1,
+ help="Number of prompts to generate.",
+ )
+ parser.add_argument(
+ "--txt-prompts",
+ type=str,
+ default=None,
+ help="Path to a .txt file with one prompt per line (preferred).",
+ )
+
+ return parser.parse_args()
+
+
+if __name__ == "__main__":
+ args = parse_args()
+ main(args)
+
+ # use examples:
+ # python end2end.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type text
+ # python end2end.py --model /custom_path/Qwen3-Omni-30B-A3B-Instruct --query-type use_video_only
diff --git a/examples/offline_inference/qwen3_omni/run_multiple_prompts.sh b/examples/offline_inference/qwen3_omni/run_multiple_prompts.sh
new file mode 100644
index 000000000..2ec8a1c57
--- /dev/null
+++ b/examples/offline_inference/qwen3_omni/run_multiple_prompts.sh
@@ -0,0 +1,3 @@
+python end2end.py --output-wav output_audio \
+ --query-type text \
+ --txt-prompts top10.txt
diff --git a/examples/offline_inference/qwen3_omni/run_single_prompt.sh b/examples/offline_inference/qwen3_omni/run_single_prompt.sh
new file mode 100644
index 000000000..c6ca09da6
--- /dev/null
+++ b/examples/offline_inference/qwen3_omni/run_single_prompt.sh
@@ -0,0 +1,2 @@
+python end2end.py --output-wav output_audio \
+ --query-type use_audio
diff --git a/examples/offline_inference/qwen_image/README.md b/examples/offline_inference/qwen_image/README.md
new file mode 100644
index 000000000..85d364568
--- /dev/null
+++ b/examples/offline_inference/qwen_image/README.md
@@ -0,0 +1,43 @@
+# Qwen-Image Offline Inference
+
+This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Image` using vLLM-Omni:
+
+- `text_to_image.py`: command-line script for single image generation.
+- `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
+
+
+## Local CLI Usage
+
+```bash
+python text_to_image.py \
+ --prompt "a cup of coffee on the table" \
+ --seed 42 \
+ --cfg_scale 4.0 \
+ --num_images_per_prompt 1 \
+ --num_inference_steps 50 \
+ --height 1024 \
+ --width 1024 \
+ --output outputs/coffee.png
+```
+
+Key arguments:
+
+- `--prompt`: text description (string).
+- `--seed`: integer seed for deterministic sampling.
+- `--cfg_scale`: true CFG scale (model-specific guidance strength).
+- `--num_images_per_prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
+- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
+- `--height/--width`: output resolution (defaults 1024x1024).
+- `--output`: path to save the generated PNG.
+
+> ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
+
+## Web UI Demo
+
+Launch the gradio demo:
+
+```bash
+python gradio_demo.py --port 7862
+```
+
+Then open `http://localhost:7862/` on your local browser to interact with the web UI.
diff --git a/examples/offline_inference/qwen_image/gradio_demo.py b/examples/offline_inference/qwen_image/gradio_demo.py
new file mode 100644
index 000000000..7906054a0
--- /dev/null
+++ b/examples/offline_inference/qwen_image/gradio_demo.py
@@ -0,0 +1,214 @@
+import argparse
+from functools import lru_cache
+
+import gradio as gr
+import torch
+
+from vllm_omni.entrypoints.omni import Omni
+
+ASPECT_RATIOS: dict[str, tuple[int, int]] = {
+ "1:1": (1328, 1328),
+ "16:9": (1664, 928),
+ "9:16": (928, 1664),
+ "4:3": (1472, 1140),
+ "3:4": (1140, 1472),
+ "3:2": (1584, 1056),
+ "2:3": (1056, 1584),
+}
+ASPECT_RATIO_CHOICES = [f"{ratio} ({w}x{h})" for ratio, (w, h) in ASPECT_RATIOS.items()]
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description="Gradio demo for Qwen-Image offline inference.")
+ parser.add_argument("--model", default="Qwen/Qwen-Image", help="Diffusion model name or local path.")
+ parser.add_argument(
+ "--height",
+ type=int,
+ default=1328,
+ help="Default image height (must match one of the supported presets).",
+ )
+ parser.add_argument(
+ "--width",
+ type=int,
+ default=1328,
+ help="Default image width (must match one of the supported presets).",
+ )
+ parser.add_argument("--default-prompt", default="a cup of coffee on the table", help="Initial prompt shown in UI.")
+ parser.add_argument("--default-seed", type=int, default=42, help="Initial seed shown in UI.")
+ parser.add_argument("--default-cfg-scale", type=float, default=4.0, help="Initial CFG scale shown in UI.")
+ parser.add_argument(
+ "--num_inference_steps",
+ type=int,
+ default=50,
+ help="Default number of denoising steps shown in the UI.",
+ )
+ parser.add_argument("--ip", default="127.0.0.1", help="Host/IP for Gradio `launch`.")
+ parser.add_argument("--port", type=int, default=7862, help="Port for Gradio `launch`.")
+ parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.")
+ args = parser.parse_args()
+ args.aspect_ratio_label = next(
+ (ratio for ratio, dims in ASPECT_RATIOS.items() if dims == (args.width, args.height)),
+ None,
+ )
+ if args.aspect_ratio_label is None:
+ supported = ", ".join(f"{ratio} ({w}x{h})" for ratio, (w, h) in ASPECT_RATIOS.items())
+ parser.error(f"Unsupported resolution {args.width}x{args.height}. Please pick one of: {supported}.")
+ return args
+
+
+@lru_cache(maxsize=1)
+def get_omni(model_name: str) -> Omni:
+ return Omni(model=model_name)
+
+
+def build_demo(args: argparse.Namespace) -> gr.Blocks:
+ device = "cuda"
+ omni = get_omni(args.model)
+
+ def run_inference(
+ prompt: str,
+ seed_value: float,
+ cfg_scale_value: float,
+ resolution_choice: str,
+ num_steps_value: float,
+ num_images_choice: float,
+ ):
+ if not prompt or not prompt.strip():
+ raise gr.Error("Please enter a non-empty prompt.")
+ ratio_label = resolution_choice.split(" ", 1)[0]
+ if ratio_label not in ASPECT_RATIOS:
+ raise gr.Error(f"Unsupported aspect ratio: {ratio_label}")
+ width, height = ASPECT_RATIOS[ratio_label]
+ try:
+ seed = int(seed_value)
+ num_steps = int(num_steps_value)
+ num_images = int(num_images_choice)
+ except (TypeError, ValueError) as exc:
+ raise gr.Error("Seed, inference steps, and number of images must be valid integers.") from exc
+ if num_steps <= 0:
+ raise gr.Error("Inference steps must be a positive integer.")
+ if num_images not in {1, 2, 3, 4}:
+ raise gr.Error("Number of images must be 1, 2, 3, or 4.")
+ generator = torch.Generator(device=device).manual_seed(seed)
+ images = omni.generate(
+ prompt.strip(),
+ height=height,
+ width=width,
+ generator=generator,
+ true_cfg_scale=float(cfg_scale_value),
+ num_inference_steps=num_steps,
+ num_images_per_prompt=num_images,
+ num_outputs_per_prompt=num_images,
+ )
+ return [img for img in images[:num_images]]
+
+ with gr.Blocks(
+ title="vLLM-Omni Web Serving Demo",
+ css="""
+ /* Left column button width */
+ .left-column button {
+ width: 100%;
+ }
+ /* Right preview area: fixed height, hide unnecessary buttons */
+ .fixed-image {
+ height: 660px;
+ display: flex;
+ flex-direction: column;
+ justify-content: center;
+ align-items: center;
+ }
+ .fixed-image .duplicate-button,
+ .fixed-image .svelte-drgfj2 {
+ display: none !important;
+ }
+ /* Gallery container: fill available space and center content */
+ #image-gallery {
+ width: 100%;
+ height: 100%;
+ display: flex;
+ align-items: center;
+ justify-content: center;
+ }
+ /* Gallery grid: center horizontally and vertically, set gap */
+ #image-gallery .grid {
+ display: flex;
+ flex-wrap: wrap;
+ justify-content: center;
+ align-items: center;
+ align-content: center;
+ gap: 16px;
+ width: 100%;
+ height: 100%;
+ }
+ /* Gallery grid items: center content */
+ #image-gallery .grid > div {
+ display: flex;
+ align-items: center;
+ justify-content: center;
+ }
+ /* Gallery images: limit max height, maintain aspect ratio */
+ .fixed-image img {
+ max-height: 660px !important;
+ width: auto !important;
+ object-fit: contain;
+ }
+ """,
+ ) as demo:
+ gr.Markdown("# vLLM-Omni Web Serving Demo")
+ gr.Markdown(f"**Model:** {args.model}")
+
+ with gr.Row():
+ with gr.Column(scale=1, elem_classes="left-column"):
+ prompt_input = gr.Textbox(
+ label="Prompt",
+ value=args.default_prompt,
+ placeholder="Describe the image you want to generate...",
+ lines=5,
+ )
+ seed_input = gr.Number(label="Seed", value=args.default_seed, precision=0)
+ cfg_input = gr.Number(label="CFG Scale", value=args.default_cfg_scale)
+ steps_input = gr.Number(
+ label="Inference Steps",
+ value=args.num_inference_steps,
+ precision=0,
+ minimum=1,
+ )
+ aspect_dropdown = gr.Dropdown(
+ label="Aspect Ratio (W:H)",
+ choices=ASPECT_RATIO_CHOICES,
+ value=f"{args.aspect_ratio_label} ({ASPECT_RATIOS[args.aspect_ratio_label][0]}x{ASPECT_RATIOS[args.aspect_ratio_label][1]})",
+ )
+ num_images = gr.Dropdown(
+ label="Number of images",
+ choices=["1", "2", "3", "4"],
+ value="1",
+ )
+ generate_btn = gr.Button("Generate", variant="primary")
+ with gr.Column(scale=2, elem_classes="fixed-image"):
+ gallery = gr.Gallery(
+ label="Preview",
+ columns=2,
+ rows=2,
+ height=660,
+ allow_preview=True,
+ show_label=True,
+ elem_id="image-gallery",
+ )
+
+ generate_btn.click(
+ fn=run_inference,
+ inputs=[prompt_input, seed_input, cfg_input, aspect_dropdown, steps_input, num_images],
+ outputs=gallery,
+ )
+
+ return demo
+
+
+def main():
+ args = parse_args()
+ demo = build_demo(args)
+ demo.launch(server_name=args.ip, server_port=args.port, share=args.share)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/offline_inference/qwen_image/text_to_image.py b/examples/offline_inference/qwen_image/text_to_image.py
new file mode 100644
index 000000000..e0a4fb8b2
--- /dev/null
+++ b/examples/offline_inference/qwen_image/text_to_image.py
@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+from pathlib import Path
+
+import torch
+
+from vllm_omni.entrypoints.omni import Omni
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description="Generate an image with Qwen-Image.")
+ parser.add_argument("--model", default="Qwen/Qwen-Image", help="Diffusion model name or local path.")
+ parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
+ parser.add_argument("--seed", type=int, default=42, help="Random seed for deterministic results.")
+ parser.add_argument(
+ "--cfg_scale",
+ type=float,
+ default=4.0,
+ help="True classifier-free guidance scale specific to Qwen-Image.",
+ )
+ parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
+ parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
+ parser.add_argument(
+ "--output",
+ type=str,
+ default="qwen_image_output.png",
+ help="Path to save the generated image (PNG).",
+ )
+ parser.add_argument(
+ "--num_images_per_prompt",
+ type=int,
+ default=1,
+ help="Number of images to generate for the given prompt.",
+ )
+ parser.add_argument(
+ "--num_inference_steps",
+ type=int,
+ default=50,
+ help="Number of denoising steps for the diffusion sampler.",
+ )
+ return parser.parse_args()
+
+
+def main():
+ args = parse_args()
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ generator = torch.Generator(device=device).manual_seed(args.seed)
+
+ omni = Omni(model=args.model)
+ images = omni.generate(
+ args.prompt,
+ height=args.height,
+ width=args.width,
+ generator=generator,
+ true_cfg_scale=args.cfg_scale,
+ num_inference_steps=args.num_inference_steps,
+ num_images_per_prompt=args.num_images_per_prompt,
+ num_outputs_per_prompt=args.num_images_per_prompt,
+ )
+
+ output_path = Path(args.output)
+ output_path.parent.mkdir(parents=True, exist_ok=True)
+ suffix = output_path.suffix or ".png"
+ stem = output_path.stem or "qwen_image_output"
+ if args.num_images_per_prompt <= 1:
+ images[0].save(output_path)
+ print(f"Saved generated image to {output_path}")
+ else:
+ for idx, img in enumerate(images):
+ save_path = output_path.parent / f"{stem}_{idx}{suffix}"
+ img.save(save_path)
+ print(f"Saved generated image to {save_path}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/online_serving/qwen2_5_omni/README.md b/examples/online_serving/qwen2_5_omni/README.md
new file mode 100644
index 000000000..da9570270
--- /dev/null
+++ b/examples/online_serving/qwen2_5_omni/README.md
@@ -0,0 +1,73 @@
+# 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 --query-type mixed_modalities
+```
+
+Send request via curl
+```bash
+bash run_curl_multimodal_generation.sh mixed_modalities
+```
+
+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
+
+## Run Local Web UI Demo
+
+This Web UI demo allows users to interact with the model through a web browser.
+
+### Running Gradio Demo
+
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+
+```bash
+python gradio_demo.py --model Qwen/Qwen2.5-Omni-7B --port 7861
+```
+
+Then open `http://localhost:7861/` on your local browser to interact with the web UI.
+
+
+### Options
+
+You can customize its basic launch parameters:
+
+```bash
+python gradio_demo.py \
+ --model Qwen/Qwen2.5-Omni-7B \
+ --ip 127.0.0.1 \
+ --port 7861 \
+ --stage-configs-path /path/to/stage_configs.yaml
+```
+
+- `--model`: Local model checkpoint to load (default `Qwen/Qwen2.5-Omni-7B`).
+- `--ip`: Host/IP for the Gradio server (default `127.0.0.1`).
+- `--port`: Port for the Gradio server (default `7861`).
+- `--stage-configs-path`: Optional path to custom stage configs YAML.
+- `--share`: Set to expose a temporary public link via Gradio.
diff --git a/examples/online_serving/gradio_demo.py b/examples/online_serving/qwen2_5_omni/gradio_demo.py
similarity index 100%
rename from examples/online_serving/gradio_demo.py
rename to examples/online_serving/qwen2_5_omni/gradio_demo.py
diff --git a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py
similarity index 100%
rename from examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py
rename to examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py
diff --git a/examples/online_serving/run_curl_multimodal_generation.sh b/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh
similarity index 100%
rename from examples/online_serving/run_curl_multimodal_generation.sh
rename to examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh
diff --git a/examples/online_serving/README.md b/examples/online_serving/qwen3_omni/README.md
similarity index 84%
rename from examples/online_serving/README.md
rename to examples/online_serving/qwen3_omni/README.md
index e93a5f249..01a0fdd28 100644
--- a/examples/online_serving/README.md
+++ b/examples/online_serving/qwen3_omni/README.md
@@ -4,16 +4,16 @@
Please refer to [README.md](../../README.md)
-## Run examples (Qwen2.5-omni)
+## Run examples (Qwen3-Omni)
Launch the server
```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --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
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
Get into the example folder
@@ -48,7 +48,7 @@ This Web UI demo allows users to interact with the model through a web browser.
Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
```bash
-python gradio_demo.py --model Qwen/Qwen2.5-Omni-7B --port 7861
+python gradio_demo.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
```
Then open `http://localhost:7861/` on your local browser to interact with the web UI.
@@ -61,7 +61,7 @@ The gradio demo also supports running with an existing API server and can be cus
```bash
python gradio_demo.py \
- --model Qwen/Qwen2.5-Omni-7B \
+ --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
--use-api-server \
--api-base http://localhost:8091/v1 \
--ip 127.0.0.1 \
diff --git a/examples/online_serving/qwen3_omni/gradio_demo.py b/examples/online_serving/qwen3_omni/gradio_demo.py
new file mode 100644
index 000000000..d6a2237e9
--- /dev/null
+++ b/examples/online_serving/qwen3_omni/gradio_demo.py
@@ -0,0 +1,537 @@
+import argparse
+import os
+import random
+import signal
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any, Optional
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+from PIL import Image
+from vllm.assets.video import video_get_metadata, video_to_ndarrays
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+
+# Import utils from offline inference example
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen3_omni"))
+try:
+ from utils import make_omni_prompt
+except ImportError:
+ # Fallback if utils doesn't exist - we'll build prompts directly
+ make_omni_prompt = None
+
+SEED = 42
+ASYNC_INIT_TIMEOUT = 600
+
+SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
+ "Qwen/Qwen3-Omni-30B-A3B-Instruct": {
+ "sampling_params": {
+ "thinker": {
+ "temperature": 0.4,
+ "top_p": 0.9,
+ "top_k": 1,
+ "max_tokens": 16384,
+ "detokenize": True,
+ "repetition_penalty": 1.05,
+ "stop_token_ids": [151645],
+ "seed": SEED,
+ },
+ "talker": {
+ "temperature": 0.0,
+ "top_p": 1.0,
+ "top_k": -1,
+ "max_tokens": 4096,
+ "detokenize": True,
+ "repetition_penalty": 1.1,
+ "stop_token_ids": [2150],
+ },
+ "code2wav": {
+ "temperature": 0.0,
+ "top_p": 1.0,
+ "top_k": -1,
+ "max_tokens": 2048,
+ "detokenize": True,
+ "repetition_penalty": 1.1,
+ },
+ },
+ },
+}
+# Ensure deterministic behavior across runs.
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+os.environ["PYTHONHASHSEED"] = str(SEED)
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Gradio demo for Qwen3-Omni online inference.")
+ parser.add_argument(
+ "--model",
+ default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
+ help="Path to local model checkpoint.",
+ )
+ parser.add_argument(
+ "--ip",
+ default="127.0.0.1",
+ help="Host/IP for gradio `launch`.",
+ )
+ parser.add_argument("--port", type=int, default=7861, help="Port for gradio `launch`.")
+ parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.")
+ parser.add_argument(
+ "--stage-configs-path",
+ type=str,
+ default=None,
+ help="Path to custom stage configs YAML file (optional).",
+ )
+ return parser.parse_args()
+
+
+def build_sampling_params(seed: int, model_key: str) -> list[SamplingParams]:
+ """Build SamplingParams objects by reusing the dict definitions."""
+ return [SamplingParams(**params_dict) for params_dict in build_sampling_params_dict(seed, model_key)]
+
+
+def build_sampling_params_dict(seed: int, model_key: str) -> list[dict]:
+ """Build sampling params as dict for HTTP API mode."""
+ model_conf = SUPPORTED_MODELS.get(model_key)
+ if model_conf is None:
+ raise ValueError(f"Unsupported model '{model_key}'")
+
+ sampling_templates: dict[str, dict[str, Any]] = model_conf["sampling_params"]
+ sampling_params: list[dict] = []
+ for stage_name, template in sampling_templates.items():
+ params = dict(template)
+ params["seed"] = seed
+ sampling_params.append(params)
+ return sampling_params
+
+
+def create_prompt_args(base_args: argparse.Namespace) -> SimpleNamespace:
+ # The prompt builder expects a minimal namespace with these attributes.
+ return SimpleNamespace(
+ model=base_args.model,
+ prompt_type="text",
+ tokenize=True,
+ use_torchvision=True,
+ legacy_omni_video=False,
+ )
+
+
+def process_audio_file(
+ audio_file: Optional[Any],
+) -> Optional[tuple[np.ndarray, int]]:
+ """Normalize Gradio audio input to (np.ndarray, sample_rate)."""
+ if audio_file is None:
+ return None
+
+ sample_rate: Optional[int] = None
+ audio_np: Optional[np.ndarray] = None
+
+ def _load_from_path(path_str: str) -> Optional[tuple[np.ndarray, int]]:
+ if not path_str:
+ return None
+ path = Path(path_str)
+ if not path.exists():
+ return None
+ data, sr = sf.read(path)
+ if data.ndim > 1:
+ data = data[:, 0]
+ return data.astype(np.float32), int(sr)
+
+ if isinstance(audio_file, tuple):
+ if len(audio_file) == 2:
+ first, second = audio_file
+ # Case 1: (sample_rate, np.ndarray)
+ if isinstance(first, (int, float)) and isinstance(second, np.ndarray):
+ sample_rate = int(first)
+ audio_np = second
+ # Case 2: (filepath, (sample_rate, np.ndarray or list))
+ elif isinstance(first, str):
+ if isinstance(second, tuple) and len(second) == 2:
+ sr_candidate, data_candidate = second
+ if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray):
+ sample_rate = int(sr_candidate)
+ audio_np = data_candidate
+ if audio_np is None:
+ loaded = _load_from_path(first)
+ if loaded is not None:
+ audio_np, sample_rate = loaded
+ # Case 3: (None, (sample_rate, np.ndarray))
+ elif first is None and isinstance(second, tuple) and len(second) == 2:
+ sr_candidate, data_candidate = second
+ if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray):
+ sample_rate = int(sr_candidate)
+ audio_np = data_candidate
+ elif len(audio_file) == 1 and isinstance(audio_file[0], str):
+ loaded = _load_from_path(audio_file[0])
+ if loaded is not None:
+ audio_np, sample_rate = loaded
+ elif isinstance(audio_file, str):
+ loaded = _load_from_path(audio_file)
+ if loaded is not None:
+ audio_np, sample_rate = loaded
+
+ if audio_np is None or sample_rate is None:
+ return None
+
+ if audio_np.ndim > 1:
+ audio_np = audio_np[:, 0]
+
+ return audio_np.astype(np.float32), sample_rate
+
+
+def process_image_file(image_file: Optional[Image.Image]) -> Optional[Image.Image]:
+ """Process image file from Gradio input.
+
+ Returns:
+ PIL Image in RGB mode or None if no image provided.
+ """
+ if image_file is None:
+ return None
+ # Convert to RGB if needed
+ if image_file.mode != "RGB":
+ image_file = image_file.convert("RGB")
+ return image_file
+
+
+def process_video_file(
+ video_file: Optional[str],
+ enable_audio_in_video: bool = False,
+ max_frames: int = 32,
+) -> Optional[tuple[np.ndarray, dict[str, Any], Optional[tuple[np.ndarray, int]]]]:
+ """Process video file and optionally extract audio track."""
+ if video_file is None:
+ return None
+
+ video_path = Path(video_file)
+ if not video_path.exists():
+ print(f"Video file not found: {video_path}")
+ return None
+
+ try:
+ frames = video_to_ndarrays(str(video_path), num_frames=max_frames)
+ metadata = video_get_metadata(str(video_path), num_frames=frames.shape[0])
+ except Exception as exc: # pylint: disable=broad-except
+ print(f"Failed to decode video {video_path}: {exc}")
+ return None
+
+ audio_tuple: Optional[tuple[np.ndarray, int]] = None
+ if enable_audio_in_video:
+ try:
+ import librosa # type: ignore import
+
+ audio_signal, sampling_rate = librosa.load(str(video_path), sr=16000)
+ audio_tuple = (audio_signal.astype(np.float32), sampling_rate)
+ except Exception as exc: # pylint: disable=broad-except
+ print(f"Failed to extract audio from video {video_path}: {exc}")
+
+ return frames, metadata, audio_tuple
+
+
+async def run_inference_async_omni_llm(
+ omni_llm: AsyncOmniLLM,
+ sampling_params: list[SamplingParams],
+ prompt_args_template: SimpleNamespace,
+ user_prompt: str,
+ audio_file: Optional[tuple[str, tuple[int, np.ndarray]]] = None,
+ image_file: Optional[Image.Image] = None,
+ video_file: Optional[str] = None,
+ use_audio_in_video: bool = False,
+):
+ """Run inference using AsyncOmniLLM directly with multimodal support."""
+ if not user_prompt.strip() and not audio_file and not image_file and not video_file:
+ return "Please provide at least a text prompt or multimodal input.", None
+
+ try:
+ # Build prompt with multimodal data
+ prompt_args = SimpleNamespace(**prompt_args_template.__dict__)
+
+ # Process multimodal inputs
+ multi_modal_data = {}
+ mm_processor_kwargs = {}
+
+ # Process audio
+ audio_data = process_audio_file(audio_file)
+ if audio_data is not None:
+ multi_modal_data["audio"] = audio_data
+
+ # Process image
+ image_data = process_image_file(image_file)
+ if image_data is not None:
+ multi_modal_data["image"] = image_data
+
+ # Process video
+ video_payload = process_video_file(video_file, enable_audio_in_video=use_audio_in_video)
+ if video_payload is not None:
+ video_frames, video_metadata, extracted_audio = video_payload
+ video_entry: Any
+ if video_metadata:
+ video_entry = (video_frames, video_metadata)
+ else:
+ video_entry = video_frames
+ multi_modal_data["video"] = video_entry
+ if use_audio_in_video and extracted_audio is not None and "audio" not in multi_modal_data:
+ multi_modal_data["audio"] = extracted_audio
+ mm_processor_kwargs["use_audio_in_video"] = True
+
+ # Build the prompt input
+ if make_omni_prompt is not None:
+ omni_prompt = make_omni_prompt(prompt_args, user_prompt)
+ # Add multimodal data if present
+ if multi_modal_data:
+ if isinstance(omni_prompt, dict):
+ omni_prompt["multi_modal_data"] = multi_modal_data
+ if mm_processor_kwargs:
+ omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
+ else:
+ # If make_omni_prompt returns a string, we need to wrap it
+ omni_prompt = {
+ "prompt": omni_prompt,
+ "multi_modal_data": multi_modal_data,
+ }
+ if mm_processor_kwargs:
+ omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
+ else:
+ # Fallback: build prompt directly
+ default_system = (
+ "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."
+ )
+ prompt = f"<|im_start|>system\n{default_system}<|im_end|>\n<|im_start|>user\n"
+ if audio_data:
+ prompt += "<|audio_bos|><|AUDIO|><|audio_eos|>"
+ if image_data:
+ prompt += "<|vision_bos|><|IMAGE|><|vision_eos|>"
+ if video_payload is not None:
+ prompt += "<|vision_bos|><|VIDEO|><|vision_eos|>"
+ if user_prompt.strip():
+ prompt += f"{user_prompt}"
+ prompt += "<|im_end|>\n<|im_start|>assistant\n"
+
+ omni_prompt = {
+ "prompt": prompt,
+ "multi_modal_data": multi_modal_data,
+ }
+ if mm_processor_kwargs:
+ omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
+
+ request_id = "0"
+ text_outputs: list[str] = []
+ audio_output = None
+
+ async for stage_outputs in omni_llm.generate(
+ prompt=omni_prompt,
+ request_id=request_id,
+ sampling_params_list=sampling_params,
+ ):
+ # stage_outputs.request_output is a RequestOutput object, not a list
+ request_output = stage_outputs.request_output
+ if stage_outputs.final_output_type == "text":
+ if request_output.outputs:
+ for output in request_output.outputs:
+ if output.text:
+ text_outputs.append(output.text)
+ elif stage_outputs.final_output_type == "audio":
+ # multimodal_output is on the RequestOutput object
+ # See vllm_omni/entrypoints/openai/serving_chat.py:680 for reference
+ if hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
+ if "audio" in request_output.multimodal_output:
+ audio_tensor = request_output.multimodal_output["audio"]
+ # Ensure audio is 1D (flatten if needed)
+ if hasattr(audio_tensor, "ndim") and audio_tensor.ndim > 1:
+ audio_tensor = audio_tensor.flatten()
+ audio_np = audio_tensor.detach().cpu().numpy()
+ audio_output = (
+ 24000, # sampling rate in Hz
+ audio_np,
+ )
+
+ text_response = "\n\n".join(text_outputs) if text_outputs else "No text output."
+ return text_response, audio_output
+ except Exception as exc: # pylint: disable=broad-except
+ return f"Inference failed: {exc}", None
+
+
+def build_interface(
+ omni_llm: AsyncOmniLLM,
+ sampling_params: list[SamplingParams],
+ prompt_args_template: SimpleNamespace,
+ model: str,
+):
+ """Build Gradio interface for AsyncOmniLLM mode."""
+
+ async def run_inference(
+ user_prompt: str,
+ audio_file: Optional[tuple[str, tuple[int, np.ndarray]]],
+ image_file: Optional[Image.Image],
+ video_file: Optional[str],
+ use_audio_in_video: bool,
+ ):
+ return await run_inference_async_omni_llm(
+ omni_llm,
+ sampling_params,
+ prompt_args_template,
+ user_prompt,
+ audio_file,
+ image_file,
+ video_file,
+ use_audio_in_video,
+ )
+
+ css = """
+ .media-input-container {
+ display: flex;
+ gap: 10px;
+ }
+ .media-input-container > div {
+ flex: 1;
+ }
+ .media-input-container .image-input,
+ .media-input-container .audio-input {
+ height: 300px;
+ }
+ .media-input-container .video-column {
+ height: 300px;
+ display: flex;
+ flex-direction: column;
+ }
+ .media-input-container .video-input {
+ flex: 1;
+ min-height: 0;
+ }
+ #generate-btn button {
+ width: 100%;
+ }
+ """
+
+ with gr.Blocks(css=css) as demo:
+ gr.Markdown("# vLLM-Omni Online Serving Demo")
+ gr.Markdown(f"**Model:** {model} \n\n")
+
+ with gr.Column():
+ with gr.Row():
+ input_box = gr.Textbox(
+ label="Text Prompt",
+ placeholder="For example: Describe what happens in the media inputs.",
+ lines=4,
+ scale=1,
+ )
+ with gr.Row(elem_classes="media-input-container"):
+ image_input = gr.Image(
+ label="Image Input (optional)",
+ type="pil",
+ sources=["upload"],
+ scale=1,
+ elem_classes="image-input",
+ )
+ with gr.Column(scale=1, elem_classes="video-column"):
+ video_input = gr.Video(
+ label="Video Input (optional)",
+ sources=["upload"],
+ elem_classes="video-input",
+ )
+ use_audio_in_video_checkbox = gr.Checkbox(
+ label="Use audio from video",
+ value=False,
+ info="Extract the video's audio track when provided.",
+ )
+ audio_input = gr.Audio(
+ label="Audio Input (optional)",
+ type="numpy",
+ sources=["upload", "microphone"],
+ scale=1,
+ elem_classes="audio-input",
+ )
+
+ with gr.Row():
+ generate_btn = gr.Button(
+ "Generate",
+ variant="primary",
+ size="lg",
+ elem_id="generate-btn",
+ )
+
+ with gr.Row():
+ text_output = gr.Textbox(label="Text Output", lines=10, scale=2)
+ audio_output = gr.Audio(label="Audio Output", interactive=False, scale=1)
+
+ generate_btn.click(
+ fn=run_inference,
+ inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox],
+ outputs=[text_output, audio_output],
+ )
+ demo.queue()
+ return demo
+
+
+def main():
+ args = parse_args()
+ omni_llm = None
+
+ model_name = "/".join(args.model.split("/")[-2:])
+ assert model_name in SUPPORTED_MODELS, (
+ f"Unsupported model '{model_name}'. Supported models: {SUPPORTED_MODELS.keys()}"
+ )
+
+ # Register signal handlers for graceful shutdown
+ def signal_handler(sig, frame):
+ print("\nReceived interrupt signal, shutting down...")
+ if omni_llm is not None:
+ try:
+ omni_llm.shutdown()
+ except Exception as e:
+ print(f"Error during shutdown: {e}")
+ sys.exit(0)
+
+ signal.signal(signal.SIGINT, signal_handler)
+ signal.signal(signal.SIGTERM, signal_handler)
+
+ print(f"Initializing AsyncOmniLLM with model: {args.model}")
+ if args.stage_configs_path:
+ print(f"Using custom stage configs: {args.stage_configs_path}")
+
+ sampling_params = build_sampling_params(SEED, model_name)
+ omni_llm = AsyncOmniLLM(
+ model=args.model,
+ stage_configs_path=args.stage_configs_path,
+ init_timeout=ASYNC_INIT_TIMEOUT,
+ )
+ print("✓ AsyncOmniLLM initialized successfully")
+ prompt_args_template = create_prompt_args(args)
+
+ demo = build_interface(
+ omni_llm,
+ sampling_params,
+ prompt_args_template,
+ args.model,
+ )
+ try:
+ demo.launch(
+ server_name=args.ip,
+ server_port=args.port,
+ share=args.share,
+ )
+ except KeyboardInterrupt:
+ print("\nShutting down...")
+ finally:
+ # Cleanup
+ if omni_llm is not None:
+ try:
+ omni_llm.shutdown()
+ except Exception as e:
+ print(f"Error during cleanup: {e}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py
new file mode 100644
index 000000000..d4522b2a4
--- /dev/null
+++ b/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py
@@ -0,0 +1,235 @@
+import base64
+from typing import NamedTuple
+
+import requests
+from openai import OpenAI
+from vllm.assets.audio import AudioAsset
+from vllm.utils import FlexibleArgumentParser
+
+# 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
+
+
+class QueryResult(NamedTuple):
+ inputs: dict
+ limit_mm_per_prompt: dict[str, int]
+
+
+def encode_base64_content_from_url(content_url: str) -> str:
+ """Encode a content retrieved from a remote url to base64 format."""
+
+ with requests.get(content_url) as response:
+ response.raise_for_status()
+ result = base64.b64encode(response.content).decode("utf-8")
+
+ return result
+
+
+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 get_text_query():
+ question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+ prompt = {
+ "role": "user",
+ "content": [
+ {
+ "type": "text",
+ "text": f"{question}",
+ }
+ ],
+ }
+ return prompt
+
+
+default_system = (
+ "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 get_video_query() -> QueryResult:
+ question = "Why is this video funny?"
+ prompt = {
+ "role": "user",
+ "content": [
+ {
+ "type": "video_url",
+ "video_url": {
+ "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+ },
+ },
+ {
+ "type": "text",
+ "text": f"{question}",
+ },
+ ],
+ }
+ return prompt
+
+
+def get_image_query() -> QueryResult:
+ question = "What is the content of this image?"
+ prompt = {
+ "role": "user",
+ "content": [
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+ },
+ },
+ {
+ "type": "text",
+ "text": f"{question}",
+ },
+ ],
+ }
+ return prompt
+
+
+def get_audio_query() -> QueryResult:
+ question = "What is the content of this audio?"
+ prompt = {
+ "role": "user",
+ "content": [
+ {
+ "type": "audio_url",
+ "audio_url": {"url": AudioAsset("mary_had_lamb").url},
+ },
+ {
+ "type": "text",
+ "text": f"{question}",
+ },
+ ],
+ }
+ return prompt
+
+
+query_map = {
+ "text": get_text_query,
+ "use_audio": get_audio_query,
+ "use_image": get_image_query,
+ "use_video": get_video_query,
+}
+
+
+def run_multimodal_generation(args) -> None:
+ model_name = args.model
+ thinker_sampling_params = {
+ "temperature": 0.4, # Deterministic
+ "top_p": 0.9,
+ "top_k": 1,
+ "max_tokens": 16384,
+ "repetition_penalty": 1.05,
+ "stop_token_ids": [151645], # Qwen EOS token <|im_end|>
+ "seed": SEED,
+ }
+
+ # Sampling parameters for Talker stage (codec generation)
+ # Stop at codec EOS token
+ talker_sampling_params = {
+ "temperature": 0.9,
+ "top_k": 50,
+ "max_tokens": 4096,
+ "seed": SEED,
+ "detokenize": False,
+ "repetition_penalty": 1.05,
+ "stop_token_ids": [2150], # TALKER_CODEC_EOS_TOKEN_ID
+ }
+
+ # # Sampling parameters for Code2Wav stage (audio generation)
+ code2wav_sampling_params = {
+ "temperature": 0.0,
+ "top_p": 1.0,
+ "top_k": -1,
+ "max_tokens": 4096 * 16,
+ "seed": SEED,
+ "detokenize": True,
+ "repetition_penalty": 1.1,
+ }
+
+ sampling_params_list = [
+ thinker_sampling_params,
+ talker_sampling_params,
+ code2wav_sampling_params,
+ ]
+
+ prompt = query_map[args.query_type]()
+ extra_body = {
+ "sampling_params_list": sampling_params_list # Optional, it has a default setting in stage_configs of the corresponding model.
+ }
+
+ if args.query_type == "use_audio_in_video":
+ extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}
+
+ chat_completion = client.chat.completions.create(
+ messages=[
+ get_system_prompt(),
+ prompt,
+ ],
+ model=model_name,
+ extra_body=extra_body,
+ )
+
+ count = 0
+ for choice in chat_completion.choices:
+ print(f"choice: {choice}")
+ 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)
+
+
+def parse_args():
+ parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+ parser.add_argument(
+ "--query-type",
+ "-q",
+ type=str,
+ default="use_video",
+ choices=query_map.keys(),
+ help="Query type.",
+ )
+ parser.add_argument(
+ "--model",
+ "-m",
+ type=str,
+ default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
+ help="Model Name / Path",
+ )
+
+ return parser.parse_args()
+
+
+if __name__ == "__main__":
+ args = parse_args()
+ run_multimodal_generation(args)
diff --git a/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh b/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh
new file mode 100644
index 000000000..95c85370c
--- /dev/null
+++ b/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh
@@ -0,0 +1,187 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Default query type
+QUERY_TYPE="${1:-mixed_modalities}"
+
+# Validate query type
+if [[ ! "$QUERY_TYPE" =~ ^(mixed_modalities|use_audio_in_video|multi_audios|text)$ ]]; then
+ echo "Error: Invalid query type '$QUERY_TYPE'"
+ echo "Usage: $0 [mixed_modalities|use_audio_in_video|multi_audios|text]"
+ echo " mixed_modalities: Audio + Image + Video + Text query"
+ echo " use_audio_in_video: Video + Text query (with audio extraction from video)"
+ echo " multi_audios: Two audio clips + Text query"
+ echo " text: Text query"
+ exit 1
+fi
+
+SEED=42
+
+thinker_sampling_params='{
+ "temperature": 0.4,
+ "top_p": 0.9,
+ "top_k": 1,
+ "max_tokens": 16384,
+ "seed": 42,
+ "detokenize": true,
+ "repetition_penalty": 1.05,
+ "stop_token_ids": [151645]
+}'
+
+talker_sampling_params='{
+ "temperature": 0.9,
+ "top_k": 50,
+ "max_tokens": 4096,
+ "seed": 42,
+ "detokenize": true,
+ "repetition_penalty": 1.05,
+ "stop_token_ids": [2150]
+}'
+
+code2wav_sampling_params='{
+ "temperature": 0.0,
+ "top_p": 1.0,
+ "top_k": -1,
+ "max_tokens": 4096 * 16,
+ "seed": 42,
+ "detokenize": true,
+ "repetition_penalty": 1.1
+}'
+# Above is optional, it has a default setting in stage_configs of the corresponding model.
+
+# Define URLs for assets
+MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
+WINNING_CALL_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/winning_call.ogg"
+CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+
+# Build user content and extra fields based on query type
+case "$QUERY_TYPE" in
+ text)
+ user_content='[
+ {
+ "type": "text",
+ "text": "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+ }
+ ]'
+ sampling_params_list='[
+ '"$thinker_sampling_params"',
+ '"$talker_sampling_params"',
+ '"$code2wav_sampling_params"'
+ ]'
+ mm_processor_kwargs="{}"
+ ;;
+ mixed_modalities)
+ user_content='[
+ {
+ "type": "audio_url",
+ "audio_url": {
+ "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
+ }
+ },
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
+ }
+ },
+ {
+ "type": "video_url",
+ "video_url": {
+ "url": "'"$SAMPLE_VIDEO_URL"'"
+ }
+ },
+ {
+ "type": "text",
+ "text": "What is recited in the audio? What is the content of this image? Why is this video funny?"
+ }
+ ]'
+ sampling_params_list='[
+ '"$thinker_sampling_params"',
+ '"$talker_sampling_params"',
+ '"$code2wav_sampling_params"'
+ ]'
+ mm_processor_kwargs="{}"
+ ;;
+ use_audio_in_video)
+ user_content='[
+ {
+ "type": "video_url",
+ "video_url": {
+ "url": "'"$SAMPLE_VIDEO_URL"'"
+ }
+ },
+ {
+ "type": "text",
+ "text": "Describe the content of the video, then convert what the baby say into text."
+ }
+ ]'
+ sampling_params_list='[
+ '"$thinker_sampling_params"',
+ '"$talker_sampling_params"',
+ '"$code2wav_sampling_params"'
+ ]'
+ mm_processor_kwargs='{
+ "use_audio_in_video": true
+ }'
+ ;;
+ multi_audios)
+ user_content='[
+ {
+ "type": "audio_url",
+ "audio_url": {
+ "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
+ }
+ },
+ {
+ "type": "audio_url",
+ "audio_url": {
+ "url": "'"$WINNING_CALL_AUDIO_URL"'"
+ }
+ },
+ {
+ "type": "text",
+ "text": "Are these two audio clips the same?"
+ }
+ ]'
+ sampling_params_list='[
+ '"$thinker_sampling_params"',
+ '"$talker_sampling_params"',
+ '"$code2wav_sampling_params"'
+ ]'
+ mm_processor_kwargs="{}"
+ ;;
+esac
+
+echo "Running query type: $QUERY_TYPE"
+echo ""
+
+
+output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
+ -H "Content-Type: application/json" \
+ -d @- <<EOF
+{
+ "model": "Qwen/Qwen2.5-Omni-7B",
+ "sampling_params_list": $sampling_params_list,
+ "mm_processor_kwargs": $mm_processor_kwargs,
+ "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": $user_content
+ }
+ ]
+}
+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')"
diff --git a/mkdocs.yml b/mkdocs.yml
index b5a54c6b7..73783ab8c 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -1,6 +1,6 @@
-site_name: vLLM-omni
+site_name: vLLM-Omni
site_description: Multi-modality models inference and serving with non-autoregressive structures
-site_author: vLLM-omni Team
+site_author: vLLM-Omni Team
site_url: https://vllm-project.github.io/vllm-omni/
repo_name: vllm-project/vllm-omni
@@ -8,30 +8,36 @@ repo_url: https://github.com/vllm-project/vllm-omni
edit_uri: edit/main/docs/
# Copyright
-copyright: Copyright © 2025 vLLM-omni Team
+copyright: Copyright © 2025 vLLM-Omni Team
# Theme
theme:
name: material
+ logo: source/logos/vllm-logo-only-light.ico
+ favicon: source/logos/vllm-logo-only-light.ico
palette:
- # Light mode
- - scheme: default
- primary: blue
- accent: blue
+ # Palette toggle for automatic mode
+ - media: "(prefers-color-scheme)"
+ toggle:
+ icon: material/brightness-auto
+ name: Switch to light mode
+ # Palette toggle for light mode
+ - media: "(prefers-color-scheme: light)"
+ scheme: default
+ primary: white
toggle:
icon: material/brightness-7
name: Switch to dark mode
- # Dark mode
- - scheme: slate
- primary: blue
- accent: blue
+ # Palette toggle for dark mode
+ - media: "(prefers-color-scheme: dark)"
+ scheme: slate
+ primary: black
toggle:
- icon: material/brightness-4
- name: Switch to light mode
+ icon: material/brightness-2
+ name: Switch to system preference
features:
- content.action.edit
- content.code.copy
- - content.tabs.link
- navigation.instant
- navigation.instant.progress
- navigation.tracking
@@ -47,17 +53,18 @@ theme:
- content.tabs
- content.tooltips
- toc.follow
- icon:
- repo: fontawesome/brands/github
+ custom_dir: docs/mkdocs/overrides
hooks:
- - docs/mkdocs/hooks/generate_examples.py
+ - docs/mkdocs/hooks/generate_api_readme.py
+ - docs/mkdocs/hooks/url_schemes.py
# Plugins
plugins:
- meta
- search
- # - awesome-nav # Uncomment if installed
+ - autorefs
+ - awesome-nav
- glightbox
- git-revision-date-localized:
# exclude files
@@ -71,6 +78,8 @@ plugins:
minify_js: true
minify_css: true
cache_safe: true
+ js_files: [docs/mkdocs/javascript/*.js]
+ css_files: [docs/mkdocs/stylesheets/*.css]
- api-autonav:
modules: ["vllm_omni"]
api_root_uri: "api"
@@ -81,33 +90,27 @@ plugins:
- 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
+ filters:
+ - "!^_" # Exclude private members (methods/classes starting with underscore)
+ summary:
+ modules: true
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
+ - 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
@@ -137,44 +140,12 @@ markdown_extensions:
# 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
+extra_javascript:
+ - mkdocs/javascript/mathjax.js
+ - https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
+ - mkdocs/javascript/edit_and_feedback.js
+ - mkdocs/javascript/slack_and_forum.js
diff --git a/pyproject.toml b/pyproject.toml
index c687f8e21..4480b54f6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,12 +5,12 @@ 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"
+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"}
+ {name = "vLLM-Omni Team", email = "vllm-project@gmail.com"}
]
keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
classifiers = [
@@ -33,6 +33,10 @@ dependencies = [
# "decord>=0.6.0", # install manually
"librosa>=0.11.0",
"resampy>=0.4.3",
+ "diffusers==0.35.2",
+ "accelerate==1.12.0",
+ "gradio==5.50",
+ "soundfile>=0.13.1"
# "vllm==0.11.0", # TODO: fix the entrypoints overwrite problem
]
@@ -46,8 +50,6 @@ dev = [
"isort>=5.12.0",
"mypy>=1.0.0",
"pre-commit>=3.0.0",
- "gradio>=5.49.1",
- "soundfile>=0.13.1"
]
docs = [
@@ -56,12 +58,13 @@ docs = [
"mkdocs-material",
"mkdocstrings-python",
"mkdocs-gen-files",
- # "mkdocs-awesome-nav", # Not available on PyPI
+ "mkdocs-awesome-nav",
"mkdocs-glightbox",
"mkdocs-git-revision-date-localized-plugin",
"mkdocs-minify-plugin",
"regex",
"pydantic",
+ "black>=23.0.0",
]
@@ -159,5 +162,9 @@ extend-ignore-identifiers-re = [
".*_thw",
".*thw",
"ein",
- ".*arange"
+ ".*arange",
+ ".*MoBA",
+ ".*temperal_downsample",
+ ".*nothink.*",
+ ".*NOTHINK.*"
]
diff --git a/tests/test_omni_llm.py b/tests/test_omni_llm.py
new file mode 100644
index 000000000..13153d6eb
--- /dev/null
+++ b/tests/test_omni_llm.py
@@ -0,0 +1,1020 @@
+import warnings
+from queue import Empty, Queue
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+# Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies.
+warnings.filterwarnings(
+ "ignore",
+ message=r"builtin type SwigPy.*has no __module__ attribute",
+ category=DeprecationWarning,
+)
+
+
+class _FakeEngineArgs(dict):
+ """Fake engine args that can be used both as object attributes and as **kwargs."""
+
+ def __init__(self, args_dict: dict[str, Any]):
+ super().__init__(args_dict)
+ # Add required attributes if not present
+ if "model_stage" not in self:
+ self["model_stage"] = None
+ if "engine_output_type" not in self:
+ self["engine_output_type"] = None
+ # Also set as attributes for object-style access
+ for key, value in self.items():
+ setattr(self, key, value)
+
+
+class _FakeStageConfig:
+ """Fake stage config object that mimics the real stage config structure."""
+
+ def __init__(self, config_dict: dict[str, Any]):
+ # engine_args needs to work both as object (for OmniStage) and as dict (for **kwargs)
+ engine_args_dict = config_dict.get("engine_args", {})
+ self.engine_args = _FakeEngineArgs(engine_args_dict)
+ self.final_output = config_dict.get("final_output", False)
+ self.final_output_type = config_dict.get("final_output_type", None)
+ self.stage_id = config_dict.get("stage_id", 0)
+ # Store original dict for reference
+ self._config_dict = config_dict
+
+
+class _FakeQueue:
+ """Fake queue using standard library Queue to replace mp.Queue."""
+
+ def __init__(self, maxsize=0):
+ self._queue = Queue(maxsize=maxsize)
+
+ def put(self, item):
+ self._queue.put(item)
+
+ def put_nowait(self, item):
+ self._queue.put_nowait(item)
+
+ def get(self):
+ return self._queue.get()
+
+ def get_nowait(self):
+ return self._queue.get_nowait()
+
+ def empty(self):
+ return self._queue.empty()
+
+
+class _FakeStage:
+ """Lightweight Stage stub for multi-process pipeline version with queue support."""
+
+ def __init__(self, config):
+ # Handle both dict and object configs
+ if isinstance(config, dict):
+ config = _FakeStageConfig(config)
+ self.config = config
+ self.stage_config = config
+ self.engine = None
+ self.engine_outputs = None
+ # Set attributes that OmniStage expects
+ self.stage_id = getattr(config, "stage_id", 0)
+ self.engine_args = config.engine_args
+ self.model_stage = getattr(config.engine_args, "model_stage", None)
+ # Allow configuring final_output and final_output_type
+ self.final_output = config.final_output if hasattr(config, "final_output") else False
+ self.final_output_type = getattr(config, "final_output_type", None)
+ # Configurable processing logic, default returns placeholder
+ processed_input = getattr(config, "_config_dict", {}).get("processed_input", ["processed"])
+ self._processed_input = processed_input
+ # Queue references (set by attach_queues)
+ self._in_q = None
+ self._out_q = None
+ self._proc = None # Mock process reference
+
+ def attach_queues(self, in_q, out_q):
+ """Attach input and output queues."""
+ self._in_q = in_q
+ self._out_q = out_q
+
+ def init_stage_worker(
+ self,
+ model: str,
+ *,
+ is_async: bool = False,
+ log_file=None,
+ shm_threshold_bytes: int = 65536,
+ ctx=None,
+ batch_timeout: int = 10,
+ ):
+ """Mock init_stage_worker: don't start real process, just send stage_ready message."""
+ # Create a mock process object
+ self._proc = MagicMock()
+ self._proc.start = MagicMock()
+ self._proc.join = MagicMock()
+ self._proc.is_alive = MagicMock(return_value=False)
+ self._proc.terminate = MagicMock()
+ # Send stage_ready message to output queue
+ if self._out_q is not None:
+ try:
+ self._out_q.put_nowait({"type": "stage_ready", "stage_id": self.stage_id})
+ except Exception:
+ pass
+
+ def stop_stage_worker(self):
+ """Mock stop_stage_worker: clean up queue references."""
+ if self._in_q is not None:
+ try:
+ self._in_q.put_nowait(None)
+ except Exception:
+ pass
+
+ def submit(self, payload: dict[str, Any]):
+ """Submit task to input queue."""
+ if self._in_q is not None:
+ self._in_q.put(payload)
+
+ def try_collect(self) -> Any:
+ """Non-blocking collect from output queue."""
+ if self._out_q is None:
+ return None
+ try:
+ return self._out_q.get_nowait()
+ except Empty:
+ return None
+
+ def set_engine_outputs(self, outputs):
+ """Set engine outputs for the stage."""
+ self.engine_outputs = outputs
+
+ def process_engine_inputs(self, stage_list, prompts):
+ """Process engine inputs: return preset processed result."""
+ return self._processed_input
+
+
+class _FakeEngine:
+ """Lightweight Engine stub: provides generate iterator output."""
+
+ def __init__(self, outputs: list[Any]):
+ self._outputs = outputs
+
+ def generate(self, prompts, sampling_params):
+ # Record the most recent prompts for outer assertions
+ self._last_prompts = prompts
+ # Simplified: return preset list at once, ensuring iterability
+ yield from self._outputs
+
+
+class _FakeStageLLM:
+ """Replace OmniStageLLM to avoid constructing real engine."""
+
+ def __init__(self, **kwargs):
+ # Allow injecting custom fake outputs, default returns single placeholder output
+ fake_outputs = kwargs.get("_fake_outputs", [[{"text": "ok"}]])
+ self._fake_engine = _FakeEngine(fake_outputs)
+
+ def generate(self, prompts, sampling_params, use_tqdm=False):
+ yield from self._fake_engine.generate(prompts, sampling_params)
+
+
+@pytest.fixture
+def fake_stage_config():
+ return {
+ # Don't include 'model' in engine_args since it's passed separately
+ "engine_args": {},
+ "final_output": True,
+ "final_output_type": "text",
+ # Second stage will use processed_input to verify the chain
+ "processed_input": ["processed-by-stage"],
+ }
+
+
+def _setup_engine_mocks(monkeypatch):
+ """Helper function to set up common engine mocks."""
+ fake_engine = MagicMock()
+ # Add necessary attributes to fake_engine
+ fake_engine.tokenizer = MagicMock()
+ fake_engine.log_stats = False
+ fake_engine.vllm_config = MagicMock()
+ fake_engine.vllm_config.model_config = MagicMock()
+ fake_engine.vllm_config.model_config.io_processor_plugin = None
+ fake_engine.get_supported_tasks = MagicMock(return_value=[])
+ fake_engine.model_config = MagicMock()
+ fake_engine.model_config.io_processor_plugin = None
+ # Add registry with resolve_model_cls method
+ fake_registry = MagicMock()
+ fake_registry.resolve_model_cls = MagicMock(return_value=(MagicMock(), "test_arch"))
+ fake_engine.model_config.registry = fake_registry
+ fake_engine.vllm_config.model_config.registry = fake_registry
+
+ monkeypatch.setattr(
+ "vllm.v1.engine.llm_engine.LLMEngine.from_engine_args",
+ lambda **kw: fake_engine,
+ raising=False,
+ )
+
+ # Mock model_config.registry.resolve_model_cls to return a tuple
+ # Use a real class instead of MagicMock to avoid inspect.getsource issues
+ class FakeModelClass:
+ pass
+
+ monkeypatch.setattr(
+ "vllm.model_executor.model_loader.utils.get_model_architecture",
+ lambda model_config: (FakeModelClass, "test_arch"),
+ raising=False,
+ )
+
+ monkeypatch.setattr(
+ "vllm.model_executor.model_loader.utils._get_model_architecture",
+ lambda model_config: (FakeModelClass, "test_arch"),
+ raising=False,
+ )
+
+ # Mock try_create_mm_pooling_model_cls to return the class as-is
+ monkeypatch.setattr(
+ "vllm.model_executor.models.adapters.try_create_mm_pooling_model_cls",
+ lambda model_cls: model_cls,
+ raising=False,
+ )
+
+ # Mock _enable_processor_cache to return False
+ monkeypatch.setattr(
+ "vllm.multimodal.cache._enable_processor_cache",
+ lambda model_config, mm_registry: False,
+ raising=False,
+ )
+
+ # Mock get_io_processor to return None
+ monkeypatch.setattr(
+ "vllm.plugins.io_processors.get_io_processor",
+ lambda vllm_config, io_processor_plugin: None,
+ raising=False,
+ )
+
+
+def _setup_multiprocessing_mocks(monkeypatch):
+ """Helper function to set up multiprocessing mocks."""
+ import multiprocessing as mp
+
+ # Mock Process
+ fake_process_class = MagicMock()
+ fake_process_instance = MagicMock()
+ fake_process_instance.start = MagicMock()
+ fake_process_instance.join = MagicMock()
+ fake_process_instance.is_alive = MagicMock(return_value=False)
+ fake_process_instance.terminate = MagicMock()
+ fake_process_class.return_value = fake_process_instance
+
+ # Mock get_context to return a context with Queue that returns _FakeQueue
+ fake_ctx = MagicMock()
+ fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize)
+ fake_ctx.Process = fake_process_class
+
+ def _mock_get_context(method):
+ return fake_ctx
+
+ monkeypatch.setattr(mp, "get_context", _mock_get_context, raising=False)
+ monkeypatch.setattr(mp, "Process", fake_process_class, raising=False)
+
+
+def _setup_ipc_mocks(monkeypatch):
+ """Helper function to set up IPC function mocks."""
+
+ # Mock _encode: simple serialization
+ def _fake_encode(obj, threshold, obj_key, shm_key):
+ return {obj_key: obj}
+
+ # Mock _load: extract object from result
+ def _fake_load(result, obj_key, shm_key):
+ return result.get(obj_key)
+
+ # Mock _set: calculate serialization size
+ def _fake_set(obj):
+ return str(obj).encode()
+
+ monkeypatch.setattr("vllm_omni.entrypoints.omni_llm._encode", _fake_encode, raising=False)
+ monkeypatch.setattr("vllm_omni.entrypoints.omni_llm._load", _fake_load, raising=False)
+ monkeypatch.setattr("vllm_omni.entrypoints.omni_llm._set", _fake_set, raising=False)
+
+
+def _setup_log_mocks(monkeypatch):
+ """Helper function to set up logging and stats mocks."""
+ # Mock init_stats_paths to return None files (no stats logging)
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.init_stats_paths",
+ lambda enable_stats, log_file: (None, None),
+ raising=False,
+ )
+ # Mock configure_orchestrator_logger to do nothing
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.configure_orchestrator_logger",
+ lambda logger, log_file: None,
+ raising=False,
+ )
+ # Mock OrchestratorMetrics to be a simple class that doesn't require file operations
+
+ class _FakeOrchestratorMetrics:
+ def __init__(self, num_stages, enable_stats, stats_file, overall_stats_file, wall_start_ts):
+ self.num_stages = num_stages
+ self.enable_stats = enable_stats
+ self.stage_first_ts = [None] * num_stages
+ self.stage_last_ts = [None] * num_stages
+ self.e2e_done = set()
+
+ def on_stage_metrics(self, stage_id, req_id, metrics):
+ pass
+
+ def on_finalize_request(self, stage_id, req_id, engine_outputs, start_ts):
+ self.e2e_done.add(req_id)
+
+ def on_forward(self, from_stage, to_stage, req_id, size_bytes, tx_ms, use_shm):
+ pass
+
+ def build_and_log_summary(self, final_stage_id):
+ return "Fake summary"
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.OrchestratorMetrics",
+ _FakeOrchestratorMetrics,
+ raising=False,
+ )
+
+
+@pytest.fixture(autouse=True)
+def mock_get_config(monkeypatch):
+ """Auto-mock get_config and related model loading functions to avoid model path validation."""
+ # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports
+ # This prevents ImportError when async_omni_llm is imported (which happens via omni_stage)
+ import sys
+
+ fake_tokenizer = MagicMock()
+ fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3])
+ fake_tokenizer.decode = MagicMock(return_value="test")
+
+ # Mock init_tokenizer_from_configs (used in async_omni_llm)
+ def _mock_init_tokenizer_from_configs(model_config=None, **kwargs):
+ return fake_tokenizer
+
+ # Strategy 1: Mock in the original location (vllm.transformers_utils.tokenizer)
+ # This works if the module hasn't been imported yet
+ monkeypatch.setattr(
+ "vllm.transformers_utils.tokenizer.init_tokenizer_from_configs",
+ _mock_init_tokenizer_from_configs,
+ raising=False,
+ )
+
+ # Strategy 2: If the module is already in sys.modules, patch it directly
+ tokenizer_module_path = "vllm.transformers_utils.tokenizer"
+ if tokenizer_module_path in sys.modules:
+ tokenizer_module = sys.modules[tokenizer_module_path]
+ setattr(tokenizer_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
+
+ # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni_llm
+ # This is because async_omni_llm imports processor.py, which imports this function at module level
+ # Mock length_from_prompt_token_ids_or_embeds (used in processor.py)
+ def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None):
+ # Return a reasonable default length
+ if prompt_token_ids is not None:
+ if isinstance(prompt_token_ids, list):
+ return len(prompt_token_ids)
+ elif hasattr(prompt_token_ids, "shape"):
+ return prompt_token_ids.shape[-1] if len(prompt_token_ids.shape) > 0 else 1
+ if prompt_embeds is not None:
+ if hasattr(prompt_embeds, "shape"):
+ return prompt_embeds.shape[-2] if len(prompt_embeds.shape) > 1 else 1
+ return 10 # Default length
+
+ # Mock in vllm.utils
+ monkeypatch.setattr(
+ "vllm.utils.length_from_prompt_token_ids_or_embeds",
+ _mock_length_from_prompt_token_ids_or_embeds,
+ raising=False,
+ )
+ # Also mock in processor module if it's imported
+ monkeypatch.setattr(
+ "vllm_omni.engine.processor.length_from_prompt_token_ids_or_embeds",
+ _mock_length_from_prompt_token_ids_or_embeds,
+ raising=False,
+ )
+ # If processor module is already imported, patch it directly
+ processor_module_path = "vllm_omni.engine.processor"
+ if processor_module_path in sys.modules:
+ processor_module = sys.modules[processor_module_path]
+ setattr(
+ processor_module, "length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds
+ )
+
+ # Strategy 3: Now mock async_omni_llm AFTER length_from_prompt_token_ids_or_embeds is mocked
+ # This prevents ImportError when async_omni_llm imports processor.py
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.async_omni_llm.init_tokenizer_from_configs",
+ _mock_init_tokenizer_from_configs,
+ raising=False,
+ )
+
+ # Strategy 4: If async_omni_llm is already imported, patch it directly
+ async_omni_llm_path = "vllm_omni.entrypoints.async_omni_llm"
+ if async_omni_llm_path in sys.modules:
+ async_omni_llm_module = sys.modules[async_omni_llm_path]
+ setattr(async_omni_llm_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
+
+ # Now mock get_config and other functions
+ fake_hf_config = MagicMock()
+ fake_hf_config.model_type = "qwen2_5_omni"
+
+ def _mock_get_config(model, **kwargs):
+ return fake_hf_config
+
+ monkeypatch.setattr(
+ "vllm.transformers_utils.config.get_config",
+ _mock_get_config,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.utils.get_config",
+ _mock_get_config,
+ raising=False,
+ )
+
+ # Mock transformers' cached_file to avoid downloading model configs
+ def _mock_cached_file(path_or_repo_id, *args, **kwargs):
+ import os
+ import tempfile
+
+ fake_config_file = os.path.join(tempfile.gettempdir(), "fake_config.json")
+ if not os.path.exists(fake_config_file):
+ with open(fake_config_file, "w") as f:
+ f.write('{"model_type": "qwen2_5_omni"}')
+ return fake_config_file
+
+ monkeypatch.setattr(
+ "transformers.utils.hub.cached_file",
+ _mock_cached_file,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "transformers.utils.hub.cached_files",
+ lambda path_or_repo_id, filenames, **kwargs: (
+ [_mock_cached_file(path_or_repo_id, filenames[0])] if filenames else None
+ ),
+ raising=False,
+ )
+
+
+def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config):
+ """Test that stage configs are auto-loaded when stage_configs_path is None."""
+
+ def _fake_loader(model: str):
+ return [
+ _FakeStageConfig(fake_stage_config),
+ _FakeStageConfig(fake_stage_config),
+ ]
+
+ # Remove modules from cache BEFORE setting mocks
+ import sys
+
+ for module_name in [
+ "vllm_omni.entrypoints.utils",
+ "vllm_omni.entrypoints.omni_llm",
+ "vllm_omni.entrypoints.omni_stage",
+ ]:
+ if module_name in sys.modules:
+ del sys.modules[module_name]
+
+ # Set up mocks
+ _setup_engine_mocks(monkeypatch)
+ _setup_multiprocessing_mocks(monkeypatch)
+ _setup_ipc_mocks(monkeypatch)
+ _setup_log_mocks(monkeypatch)
+
+ # Mock remove_old_logs to avoid bug (called before stage_list is created)
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.remove_old_logs",
+ lambda log_file, num_stages: None,
+ raising=False,
+ )
+
+ # Mock load_stage_configs_from_model
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+
+ # Replace OmniStage
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_stage.OmniStage",
+ lambda cfg: _FakeStage(cfg),
+ raising=False,
+ )
+
+ # Import the module after mocks are set
+ import vllm_omni.entrypoints.omni_llm as omni_llm_module
+
+ # Patch the imported function and class in the module
+ monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
+ monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
+
+ from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+ llm = OmniLLM(model="any", init_timeout=1)
+ # Verify: auto-loaded stage_configs and stage_list have consistent count
+ assert isinstance(llm.stage_configs, list)
+ assert len(llm.stage_configs) == 2
+ assert len(llm.stage_list) == 2
+ # Verify: each Stage is _FakeStage instance
+ for st in llm.stage_list:
+ assert isinstance(st, _FakeStage)
+ # Verify: queues are attached
+ for st in llm.stage_list:
+ assert st._in_q is not None
+ assert st._out_q is not None
+ # Verify: all stages are ready
+ assert len(llm._stages_ready) == 2
+
+
+def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
+ """Test that generate raises ValueError when sampling_params_list length doesn't match."""
+
+ def _fake_loader(model: str):
+ return [_FakeStageConfig(fake_stage_config)]
+
+ import sys
+
+ for module_name in [
+ "vllm_omni.entrypoints.utils",
+ "vllm_omni.entrypoints.omni_llm",
+ "vllm_omni.entrypoints.omni_stage",
+ ]:
+ if module_name in sys.modules:
+ del sys.modules[module_name]
+
+ _setup_engine_mocks(monkeypatch)
+ _setup_multiprocessing_mocks(monkeypatch)
+ _setup_ipc_mocks(monkeypatch)
+ _setup_log_mocks(monkeypatch)
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.remove_old_logs",
+ lambda log_file, num_stages: None,
+ raising=False,
+ )
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_stage.OmniStage",
+ lambda cfg: _FakeStage(cfg),
+ raising=False,
+ )
+
+ import vllm_omni.entrypoints.omni_llm as omni_llm_module
+
+ monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
+ monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
+
+ from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+ llm = OmniLLM(model="any", init_timeout=1)
+ with pytest.raises(ValueError):
+ llm.generate(prompts=["hi"], sampling_params_list=[])
+
+
+def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
+ """Test multi-stage generation pipeline with queue polling."""
+ stage_cfg0 = dict(fake_stage_config)
+ stage_cfg1 = dict(fake_stage_config)
+ stage_cfg1["processed_input"] = ["processed-for-stage-1"]
+
+ def _fake_loader(model: str):
+ return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+
+ import sys
+
+ for module_name in [
+ "vllm_omni.entrypoints.utils",
+ "vllm_omni.entrypoints.omni_llm",
+ "vllm_omni.entrypoints.omni_stage",
+ ]:
+ if module_name in sys.modules:
+ del sys.modules[module_name]
+
+ _setup_engine_mocks(monkeypatch)
+ _setup_multiprocessing_mocks(monkeypatch)
+ _setup_ipc_mocks(monkeypatch)
+ _setup_log_mocks(monkeypatch)
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.remove_old_logs",
+ lambda log_file, num_stages: None,
+ raising=False,
+ )
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_stage.OmniStage",
+ lambda cfg: _FakeStage(cfg),
+ raising=False,
+ )
+
+ import vllm_omni.entrypoints.omni_llm as omni_llm_module
+
+ monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
+ monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
+
+ from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+ llm = OmniLLM(model="any", init_timeout=1)
+
+ # Simulate worker behavior: manually put results into output queues
+ # Note: We put results before calling generate, which simulates worker processes
+ # that have already completed. The polling loop will collect them in stage order.
+ # Stage 0 output (will be collected first)
+ llm.stage_list[0]._out_q.put_nowait(
+ {
+ "request_id": 0,
+ "engine_outputs": [{"stage": 0, "text": "s0"}],
+ "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+ }
+ )
+ # Stage 1 output (will be collected after stage 0 forwards to it)
+ # Note: In real flow, stage 1 result would appear after stage 0 forwards,
+ # but for testing we pre-populate it. The polling loop processes stages
+ # in order, so stage 0 result will be collected first, then forwarded,
+ # then stage 1 result will be collected.
+ llm.stage_list[1]._out_q.put_nowait(
+ {
+ "request_id": 0,
+ "engine_outputs": [{"stage": 1, "text": "s1"}],
+ "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+ }
+ )
+
+ sampling_params_list = [object(), object()]
+ prompts = ["hi"]
+ outputs = llm.generate(prompts=prompts, sampling_params_list=sampling_params_list)
+
+ # Both stages have final_output=True, so should aggregate two OmniRequestOutput
+ assert len(outputs) == 2
+ # Verify stage outputs are set
+ assert llm.stage_list[0].engine_outputs == [{"stage": 0, "text": "s0"}]
+ assert llm.stage_list[1].engine_outputs == [{"stage": 1, "text": "s1"}]
+ # Verify stage 0 input queue received the task
+ assert not llm.stage_list[0]._in_q.empty()
+ # Verify stage 1 received forwarded task (process_engine_inputs was called)
+ assert llm.stage_list[1].process_engine_inputs([], []) is not None
+
+
+def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
+ """Test that generate returns empty list when all stages have final_output=False."""
+ stage_cfg0 = dict(fake_stage_config)
+ stage_cfg1 = dict(fake_stage_config)
+ stage_cfg0["final_output"] = False
+ stage_cfg1["final_output"] = False
+
+ def _fake_loader(model: str):
+ return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+
+ import sys
+
+ for module_name in [
+ "vllm_omni.entrypoints.utils",
+ "vllm_omni.entrypoints.omni_llm",
+ "vllm_omni.entrypoints.omni_stage",
+ ]:
+ if module_name in sys.modules:
+ del sys.modules[module_name]
+
+ _setup_engine_mocks(monkeypatch)
+ _setup_multiprocessing_mocks(monkeypatch)
+ _setup_ipc_mocks(monkeypatch)
+ _setup_log_mocks(monkeypatch)
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.remove_old_logs",
+ lambda log_file, num_stages: None,
+ raising=False,
+ )
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_stage.OmniStage",
+ lambda cfg: _FakeStage(cfg),
+ raising=False,
+ )
+
+ import vllm_omni.entrypoints.omni_llm as omni_llm_module
+
+ monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
+ monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
+
+ from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+ llm = OmniLLM(model="any", init_timeout=1)
+
+ # Simulate worker behavior: put results into output queues
+ llm.stage_list[0]._out_q.put_nowait(
+ {
+ "request_id": 0,
+ "engine_outputs": [{"stage": 0}],
+ "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+ }
+ )
+ llm.stage_list[1]._out_q.put_nowait(
+ {
+ "request_id": 0,
+ "engine_outputs": [{"stage": 1}],
+ "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+ }
+ )
+
+ outputs = llm.generate(prompts=["p"], sampling_params_list=[object(), object()])
+ assert outputs == []
+
+
+def test_generate_sampling_params_none_raises(monkeypatch, fake_stage_config):
+ """Test that generate raises ValueError when sampling_params_list is None."""
+
+ def _fake_loader(model: str):
+ return [_FakeStageConfig(fake_stage_config)]
+
+ import sys
+
+ for module_name in [
+ "vllm_omni.entrypoints.utils",
+ "vllm_omni.entrypoints.omni_llm",
+ "vllm_omni.entrypoints.omni_stage",
+ ]:
+ if module_name in sys.modules:
+ del sys.modules[module_name]
+
+ _setup_engine_mocks(monkeypatch)
+ _setup_multiprocessing_mocks(monkeypatch)
+ _setup_ipc_mocks(monkeypatch)
+ _setup_log_mocks(monkeypatch)
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.remove_old_logs",
+ lambda log_file, num_stages: None,
+ raising=False,
+ )
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_stage.OmniStage",
+ lambda cfg: _FakeStage(cfg),
+ raising=False,
+ )
+
+ import vllm_omni.entrypoints.omni_llm as omni_llm_module
+
+ monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
+ monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
+
+ from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+ llm = OmniLLM(model="any", init_timeout=1)
+ with pytest.raises(ValueError):
+ llm.generate(prompts=["p"], sampling_params_list=None)
+
+
+def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
+ """Test that _wait_for_stages_ready handles timeout correctly."""
+
+ def _fake_loader(model: str):
+ return [_FakeStageConfig(fake_stage_config)]
+
+ import sys
+
+ for module_name in [
+ "vllm_omni.entrypoints.utils",
+ "vllm_omni.entrypoints.omni_llm",
+ "vllm_omni.entrypoints.omni_stage",
+ ]:
+ if module_name in sys.modules:
+ del sys.modules[module_name]
+
+ _setup_engine_mocks(monkeypatch)
+ _setup_multiprocessing_mocks(monkeypatch)
+ _setup_ipc_mocks(monkeypatch)
+ _setup_log_mocks(monkeypatch)
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.remove_old_logs",
+ lambda log_file, num_stages: None,
+ raising=False,
+ )
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+
+ # Create a stage that doesn't send stage_ready message
+ class _FakeStageNoReady(_FakeStage):
+ def init_stage_worker(self, *args, **kwargs):
+ # Don't send stage_ready message
+ self._proc = MagicMock()
+ self._proc.start = MagicMock()
+ self._proc.join = MagicMock()
+ self._proc.is_alive = MagicMock(return_value=False)
+ self._proc.terminate = MagicMock()
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_stage.OmniStage",
+ lambda cfg: _FakeStageNoReady(cfg),
+ raising=False,
+ )
+
+ import vllm_omni.entrypoints.omni_llm as omni_llm_module
+
+ monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
+ monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStageNoReady(cfg))
+
+ from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+ # Use very short timeout
+ llm = OmniLLM(model="any", init_timeout=0.01)
+ # Verify that no stages are ready
+ assert len(llm._stages_ready) == 0
+
+
+def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
+ """Test that generate handles error messages from stages correctly."""
+
+ def _fake_loader(model: str):
+ return [_FakeStageConfig(fake_stage_config)]
+
+ import sys
+
+ for module_name in [
+ "vllm_omni.entrypoints.utils",
+ "vllm_omni.entrypoints.omni_llm",
+ "vllm_omni.entrypoints.omni_stage",
+ ]:
+ if module_name in sys.modules:
+ del sys.modules[module_name]
+
+ _setup_engine_mocks(monkeypatch)
+ _setup_multiprocessing_mocks(monkeypatch)
+ _setup_ipc_mocks(monkeypatch)
+ _setup_log_mocks(monkeypatch)
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.remove_old_logs",
+ lambda log_file, num_stages: None,
+ raising=False,
+ )
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_stage.OmniStage",
+ lambda cfg: _FakeStage(cfg),
+ raising=False,
+ )
+
+ import vllm_omni.entrypoints.omni_llm as omni_llm_module
+
+ monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
+ monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
+
+ from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+ llm = OmniLLM(model="any", init_timeout=1)
+
+ # Put error message in output queue
+ llm.stage_list[0]._out_q.put_nowait(
+ {
+ "request_id": 0,
+ "error": "test error",
+ }
+ )
+ # Also put a valid result after error to allow the loop to complete
+ # (error handling continues the loop, so we need a valid result to finish)
+ llm.stage_list[0]._out_q.put_nowait(
+ {
+ "request_id": 0,
+ "engine_outputs": [{"stage": 0, "text": "result"}],
+ "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+ }
+ )
+
+ # Generate should handle error gracefully (log but continue)
+ sampling_params_list = [object()]
+ outputs = llm.generate(prompts=["hi"], sampling_params_list=sampling_params_list)
+ # Should return final output (error was logged but didn't stop processing)
+ assert isinstance(outputs, list)
+ # Since final_output=True, should have one output
+ assert len(outputs) == 1
+
+
+def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
+ """Test that close() sends shutdown signal to all input queues."""
+
+ def _fake_loader(model: str):
+ return [_FakeStageConfig(fake_stage_config)]
+
+ import sys
+
+ for module_name in [
+ "vllm_omni.entrypoints.utils",
+ "vllm_omni.entrypoints.omni_llm",
+ "vllm_omni.entrypoints.omni_stage",
+ ]:
+ if module_name in sys.modules:
+ del sys.modules[module_name]
+
+ _setup_engine_mocks(monkeypatch)
+ _setup_multiprocessing_mocks(monkeypatch)
+ _setup_ipc_mocks(monkeypatch)
+ _setup_log_mocks(monkeypatch)
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.remove_old_logs",
+ lambda log_file, num_stages: None,
+ raising=False,
+ )
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_llm.load_stage_configs_from_model",
+ _fake_loader,
+ raising=False,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.omni_stage.OmniStage",
+ lambda cfg: _FakeStage(cfg),
+ raising=False,
+ )
+
+ import vllm_omni.entrypoints.omni_llm as omni_llm_module
+
+ monkeypatch.setattr(omni_llm_module, "load_stage_configs_from_model", _fake_loader)
+ monkeypatch.setattr(omni_llm_module, "OmniStage", lambda cfg: _FakeStage(cfg))
+
+ from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+ llm = OmniLLM(model="any", init_timeout=1)
+
+ # Call close
+ llm.close()
+
+ # Verify shutdown signal (None) was sent to input queue
+ # Use get_nowait to avoid blocking (close() uses put_nowait, so should be safe)
+ try:
+ shutdown_signal = llm.stage_list[0]._in_q.get_nowait()
+ assert shutdown_signal is None
+ except Empty:
+ # If queue was already empty or only had stage_ready, that's also acceptable
+ # The important thing is that close() was called without error
+ pass
+
+ # Verify stop_stage_worker was called (process should be set)
+ assert llm.stage_list[0]._proc is not None
diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py
index 2e742a027..89c7cbfc6 100644
--- a/vllm_omni/__init__.py
+++ b/vllm_omni/__init__.py
@@ -1,5 +1,5 @@
"""
-vLLM-omni: Multi-modality models inference and serving with
+vLLM-Omni: Multi-modality models inference and serving with
non-autoregressive structures.
This package extends vLLM beyond traditional text-based, autoregressive
@@ -19,20 +19,21 @@ except ModuleNotFoundError as exc: # pragma: no cover - optional dependency
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 .entrypoints.omni import Omni
-# from .version import __version__, __version_tuple__ # isort:skip
+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
+ "__version__",
+ "__version_tuple__",
# Main components
- "OmniLLM",
+ # "Omni",
"AsyncOmniLLM",
# Configuration
"OmniModelConfig",
diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py
index 0fae07e7a..cc7e8132e 100644
--- a/vllm_omni/config/__init__.py
+++ b/vllm_omni/config/__init__.py
@@ -1,5 +1,5 @@
"""
-Configuration module for vLLM-omni.
+Configuration module for vLLM-Omni.
"""
from vllm_omni.config.model import OmniModelConfig
diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py
index 67bcc5fad..ef9710e90 100644
--- a/vllm_omni/config/model.py
+++ b/vllm_omni/config/model.py
@@ -1,21 +1,72 @@
-from typing import Optional
+import json
+import warnings
+from importlib.util import find_spec
+from typing import Any, Literal, Optional
+import torch
+import vllm.envs as envs
+from omegaconf import DictConfig, OmegaConf
from pydantic import ConfigDict
from pydantic.dataclasses import dataclass
from vllm.config import ModelConfig, config
+from vllm.config.model import (
+ _RUNNER_CONVERTS,
+ _RUNNER_TASKS,
+ ConvertType,
+ RunnerOption,
+ TaskOption,
+ _get_and_verify_dtype,
+ get_served_model_name,
+)
+from vllm.config.multimodal import MMCacheType, MMEncoderTPMode, MultiModalConfig
+from vllm.config.pooler import PoolerConfig
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.transformers_utils.config import (
+ get_config,
+ get_hf_image_processor_config,
+ get_hf_text_config,
+ get_pooling_config,
+ is_interleaved,
+)
+from vllm.transformers_utils.utils import maybe_model_redirect
import vllm_omni.model_executor.models as me_models
+logger = init_logger(__name__)
+
@config
@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
class OmniModelConfig(ModelConfig):
- """Configuration for Omni models, extending the base ModelConfig."""
+ """Configuration for Omni models, extending the base ModelConfig.
+
+ This configuration class extends the base vLLM ModelConfig with
+ omni-specific fields for multi-stage pipeline processing.
+
+ Attributes:
+ stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
+ model_stage: Stage type identifier, e.g., "thinker" or "talker"
+ (default: "thinker")
+ model_arch: Model architecture name
+ (default: "Qwen2_5OmniForConditionalGeneration")
+ engine_output_type: Optional output type specification for the engine.
+ Used to route outputs to appropriate processors (e.g., "image",
+ "audio", "latents"). If None, output type is inferred.
+
+ Example:
+ >>> config = OmniModelConfig(
+ ... stage_id=0,
+ ... model_stage="thinker",
+ ... model_arch="Qwen2_5OmniForConditionalGeneration"
+ ... )
+ """
stage_id: int = 0
model_stage: str = "thinker"
model_arch: str = "Qwen2_5OmniForConditionalGeneration"
engine_output_type: Optional[str] = None
+ hf_config_name: Optional[str] = None
@property
def registry(self):
@@ -24,3 +75,313 @@ class OmniModelConfig(ModelConfig):
@property
def architectures(self) -> list[str]:
return [self.model_arch]
+
+ def draw_hf_text_config(self):
+ # transformers' get_text_config method is used to get the text config from thinker_config.
+ # to handle the case that each model stage has their own text config,
+ # we need to draw the text config from the corresponding model stage.
+ if self.hf_config_name is None:
+ return get_hf_text_config(self.hf_config)
+ return getattr(self.hf_config, self.hf_config_name).get_text_config()
+
+ def __post_init__(
+ self,
+ # Multimodal config init vars
+ limit_mm_per_prompt: Optional[dict[str, int]],
+ media_io_kwargs: Optional[dict[str, dict[str, Any]]],
+ mm_processor_kwargs: Optional[dict[str, Any]],
+ mm_processor_cache_gb: Optional[float],
+ mm_processor_cache_type: Optional[MMCacheType],
+ mm_shm_cache_max_object_size_mb: Optional[int],
+ mm_encoder_tp_mode: Optional[MMEncoderTPMode],
+ interleave_mm_strings: Optional[bool],
+ skip_mm_profiling: Optional[bool],
+ video_pruning_rate: Optional[float],
+ ) -> None:
+ # Set the default seed to 0 in V1.
+ if envs.VLLM_USE_V1 and self.seed is None:
+ self.seed = 0
+ if not envs.VLLM_ENABLE_V1_MULTIPROCESSING:
+ logger.warning(
+ "The global random seed is set to %d. Since "
+ "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may "
+ "affect the random state of the Python process that "
+ "launched vLLM.",
+ self.seed,
+ )
+
+ # Keep set served_model_name before maybe_model_redirect(self.model)
+ self.served_model_name = get_served_model_name(self.model, self.served_model_name)
+ self.model = maybe_model_redirect(self.model)
+ # The tokenizer is consistent with the model by default.
+ if self.tokenizer is None:
+ self.tokenizer = self.model
+ if self.tokenizer_revision is None:
+ self.tokenizer_revision = self.revision
+ self.tokenizer = maybe_model_redirect(self.tokenizer)
+
+ if isinstance(self.hf_config_path, str):
+ self.hf_config_path = maybe_model_redirect(self.hf_config_path)
+
+ if callable(self.hf_overrides):
+ hf_overrides_kw: dict[str, Any] = {}
+ hf_overrides_fn = self.hf_overrides
+ dict_overrides: dict[str, Any] = {}
+ else:
+ # Separate dict overrides from flat ones; convert OmegaConf to dict
+ hf_overrides_kw = {}
+ dict_overrides = {}
+ for key, value in self.hf_overrides.items():
+ if isinstance(value, DictConfig):
+ value = OmegaConf.to_container(value, resolve=True, structured_config_mode=None)
+ if isinstance(value, dict):
+ dict_overrides[key] = value
+ else:
+ hf_overrides_kw[key] = value
+ hf_overrides_fn = None
+
+ if self.rope_scaling:
+ hf_override: dict[str, Any] = {"rope_scaling": self.rope_scaling}
+ hf_overrides_kw.update(hf_override)
+ hf_overrides_str = json.dumps(hf_overrides_kw)
+ msg = (
+ "`--rope-scaling` will be removed in a future release. "
+ f"'Please instead use `--hf-overrides '{hf_overrides_str}'`"
+ )
+ warnings.warn(DeprecationWarning(msg), stacklevel=2)
+ if self.rope_theta is not None:
+ hf_override = {"rope_theta": self.rope_theta}
+ hf_overrides_kw.update(hf_override)
+ hf_overrides_str = json.dumps(hf_overrides_kw)
+ msg = (
+ "`--rope-theta` will be removed in a future release. "
+ f"'Please instead use `--hf-overrides '{hf_overrides_str}'`"
+ )
+ warnings.warn(DeprecationWarning(msg), stacklevel=2)
+
+ self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer)
+
+ if (backend := envs.VLLM_ATTENTION_BACKEND) and backend == "FLASHINFER" and find_spec("flashinfer") is None:
+ raise ValueError(
+ "VLLM_ATTENTION_BACKEND is set to FLASHINFER, but flashinfer "
+ "module was not found. See "
+ "https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile "
+ "for instructions on how to install it."
+ )
+
+ if self.override_attention_dtype is not None and not current_platform.is_rocm():
+ warnings.warn("override-attention-dtype is set but not using ROCm platform", stacklevel=2)
+
+ if self.enable_sleep_mode and not current_platform.is_sleep_mode_available():
+ raise ValueError("Sleep mode is not supported on current platform.")
+
+ hf_config = get_config(
+ self.hf_config_path or self.model,
+ self.trust_remote_code,
+ self.revision,
+ self.code_revision,
+ self.config_format,
+ hf_overrides_kw=hf_overrides_kw,
+ hf_overrides_fn=hf_overrides_fn,
+ )
+
+ if dict_overrides:
+ self._apply_dict_overrides(hf_config, dict_overrides)
+
+ self.hf_config = hf_config
+ self.hf_text_config = self.draw_hf_text_config()
+ self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None)
+ self.encoder_config = self._get_encoder_config()
+ self.hf_image_processor_config = get_hf_image_processor_config(
+ self.model, hf_token=self.hf_token, revision=self.revision
+ )
+
+ architectures = self.architectures
+ registry = self.registry
+ is_generative_model = registry.is_text_generation_model(architectures, self)
+ is_pooling_model = registry.is_pooling_model(architectures, self)
+
+ def _task_to_convert(task: TaskOption) -> ConvertType:
+ if task == "embedding" or task == "embed":
+ return "embed"
+ if task == "classify":
+ return "classify"
+ if task == "reward":
+ return "reward"
+ if task == "score":
+ new_task = self._get_default_pooling_task(architectures)
+ return "classify" if new_task == "classify" else "embed"
+
+ return "none"
+
+ if self.task is not None:
+ runner: RunnerOption = "auto"
+ convert: ConvertType | Literal["auto"] = "auto" # type: ignore[name-defined]
+ msg_prefix = (
+ "The 'task' option has been deprecated and will be removed in v0.13.0 or v1.0, whichever comes first."
+ )
+ msg_hint = "Please remove this option."
+
+ is_generative_task = self.task in _RUNNER_TASKS["generate"]
+ is_pooling_task = self.task in _RUNNER_TASKS["pooling"]
+
+ if is_generative_model and is_pooling_model:
+ if is_generative_task:
+ runner = "generate"
+ convert = "auto" # type: ignore[assignment]
+ msg_hint = (
+ "Please replace this option with `--runner "
+ "generate` to continue using this model "
+ "as a generative model."
+ )
+ elif is_pooling_task:
+ runner = "pooling"
+ convert = "auto" # type: ignore[assignment]
+ msg_hint = (
+ "Please replace this option with `--runner "
+ "pooling` to continue using this model "
+ "as a pooling model."
+ )
+ else: # task == "auto"
+ pass
+ elif is_generative_model or is_pooling_model:
+ if is_generative_task:
+ runner = "generate"
+ convert = "auto" # type: ignore[assignment]
+ msg_hint = "Please remove this option"
+ elif is_pooling_task:
+ runner = "pooling"
+ convert = _task_to_convert(self.task)
+ msg_hint = (
+ "Please replace this option with `--convert "
+ f"{convert}` to continue using this model "
+ "as a pooling model."
+ )
+ else: # task == "auto"
+ pass
+ else:
+ raise AssertionError(
+ f"The model should be a generative or pooling model when task is set to {self.task!r}."
+ )
+
+ self.runner = runner
+ self.convert = convert # type: ignore[assignment]
+
+ msg = f"{msg_prefix} {msg_hint}"
+ warnings.warn(msg, DeprecationWarning, stacklevel=2)
+
+ self.runner_type = self._get_runner_type(architectures, self.runner)
+ self.convert_type = self._get_convert_type(architectures, self.runner_type, self.convert)
+
+ if self.runner_type == "generate" and not is_generative_model:
+ generate_converts = _RUNNER_CONVERTS["generate"]
+ if self.convert_type not in generate_converts:
+ raise ValueError("This model does not support `--runner generate`.")
+ if self.runner_type == "pooling" and not is_pooling_model:
+ pooling_converts = _RUNNER_CONVERTS["pooling"]
+ if self.convert_type not in pooling_converts:
+ convert_option = "<" + "|".join(pooling_converts) + ">"
+ raise ValueError(
+ "This model does not support `--runner pooling`. "
+ f"You can pass `--convert {convert_option} to adapt "
+ "it into a pooling model."
+ )
+
+ # Note: Initialize these attributes early because transformers fallback
+ # may fail to load dynamic modules in child processes
+ model_info, arch = registry.inspect_model_cls(architectures, self)
+ self._model_info = model_info
+ self._architecture = arch
+ logger.info("Resolved architecture: %s", arch)
+
+ # Init pooler config if needed
+ if self.runner_type == "pooling":
+ if self.override_pooler_config is not None:
+ logger.warning_once(
+ "`override_pooler_config` is deprecated and will be "
+ "removed in v0.12.0 or v1.0.0, whichever is sooner. "
+ "Please use `pooler_config` instead."
+ )
+
+ if isinstance(self.override_pooler_config, dict):
+ self.pooler_config = PoolerConfig(**self.override_pooler_config)
+ else:
+ self.pooler_config = self.override_pooler_config
+
+ if self.pooler_config is None:
+ self.pooler_config = PoolerConfig()
+
+ base_config = get_pooling_config(self.model, self.revision)
+ if base_config is not None:
+ for k, v in base_config.items():
+ if getattr(self.pooler_config, k) is None:
+ setattr(self.pooler_config, k, v)
+
+ default_pooling_type = self._model_info.default_pooling_type
+ if self.pooler_config.pooling_type is None:
+ self.pooler_config.pooling_type = default_pooling_type
+
+ self.dtype: torch.dtype = _get_and_verify_dtype(
+ self.model,
+ self.hf_config,
+ self.dtype,
+ is_pooling_model=self.runner_type == "pooling",
+ revision=self.revision,
+ )
+
+ # Interleaved attention is not supported by some backends in V0
+ if (
+ not self.disable_sliding_window
+ and is_interleaved(self.hf_text_config)
+ and not envs.VLLM_USE_V1
+ and (backend := envs.VLLM_ATTENTION_BACKEND) in ("XFORMERS", "FLASHINFER")
+ ):
+ logger.warning_once(
+ "%s has interleaved attention, which is currently not "
+ "supported by the %s backend. Disabling sliding window and "
+ "capping the max length to the sliding window size (%d).",
+ self.hf_text_config.model_type,
+ backend,
+ self.hf_text_config.sliding_window,
+ )
+ self.disable_sliding_window = True
+
+ self.original_max_model_len = self.max_model_len
+ self.max_model_len = self.get_and_verify_max_len(self.max_model_len)
+ # Init multimodal config if needed
+ if self._model_info.supports_multimodal:
+ if mm_encoder_tp_mode == "data" and not self._model_info.supports_multimodal_encoder_tp_data:
+ logger.warning_once(
+ "This model does not support `--mm-encoder-tp-mode data`. "
+ "Falling back to `--mm-encoder-tp-mode weights`."
+ )
+ mm_encoder_tp_mode = "weights"
+
+ mm_config_kwargs = dict(
+ limit_per_prompt=limit_mm_per_prompt,
+ media_io_kwargs=media_io_kwargs,
+ mm_processor_kwargs=mm_processor_kwargs,
+ mm_processor_cache_gb=mm_processor_cache_gb,
+ mm_processor_cache_type=mm_processor_cache_type,
+ mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
+ mm_encoder_tp_mode=mm_encoder_tp_mode,
+ interleave_mm_strings=interleave_mm_strings,
+ skip_mm_profiling=skip_mm_profiling,
+ video_pruning_rate=video_pruning_rate,
+ )
+
+ mm_config_kwargs = {k: v for k, v in mm_config_kwargs.items() if v is not None}
+
+ self.multimodal_config = MultiModalConfig(**mm_config_kwargs)
+
+ if self.disable_sliding_window:
+ self.hf_text_config.sliding_window = None
+
+ if not self.skip_tokenizer_init:
+ self._verify_tokenizer_mode()
+
+ self.config_updated = False
+
+ self._verify_quantization()
+ self._verify_cuda_graph()
+ self._verify_bnb_config()
diff --git a/vllm_omni/core/dit_cache_manager.py b/vllm_omni/core/dit_cache_manager.py
index 9009638d5..e07c4a167 100644
--- a/vllm_omni/core/dit_cache_manager.py
+++ b/vllm_omni/core/dit_cache_manager.py
@@ -1,5 +1,5 @@
"""
-DiT Cache Manager for vLLM-omni.
+DiT Cache Manager for vLLM-Omni.
This module provides caching functionality for DiT (Diffusion Transformer) models
to optimize inference performance and memory usage.
diff --git a/vllm_omni/core/sched/__init__.py b/vllm_omni/core/sched/__init__.py
index 4f2b89170..9b3a34ca5 100644
--- a/vllm_omni/core/sched/__init__.py
+++ b/vllm_omni/core/sched/__init__.py
@@ -1,5 +1,5 @@
"""
-Scheduling components for vLLM-omni.
+Scheduling components for vLLM-Omni.
"""
from .diffusion_scheduler import DiffusionScheduler
diff --git a/vllm_omni/core/sched/diffusion_scheduler.py b/vllm_omni/core/sched/diffusion_scheduler.py
index 0474f83ab..e8707aabf 100644
--- a/vllm_omni/core/sched/diffusion_scheduler.py
+++ b/vllm_omni/core/sched/diffusion_scheduler.py
@@ -14,12 +14,22 @@ from vllm_omni.outputs import OmniModelRunnerOutput
class DiffusionScheduler(OmniScheduler):
+ """Scheduler for diffusion-based non-autoregressive stages.
+
+ Implements a fast-path scheduling strategy for diffusion models where
+ all input tokens are fed at once. Falls back to standard vLLM scheduling
+ if the token budget cannot accommodate all tokens in a single iteration.
+ """
+
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.
+ """Schedule requests for diffusion model execution.
+
+ Uses a fast-path approach that feeds all input tokens at once
+ (or 1 placeholder token if input is empty). Falls back to
+ standard scheduling if token budget is insufficient.
+
+ Returns:
+ SchedulerOutput containing scheduled requests and token allocations
"""
# Select requests with zero prompt and using pooling
@@ -239,6 +249,12 @@ class DiffusionScheduler(OmniScheduler):
# Get prompt logprobs for this request.
prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
+ # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput
+ pooling_output_dict = None
+ if pooler_output is not None:
+ # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type
+ # which expects Optional[dict[str, torch.Tensor]]
+ pooling_output_dict = {"model_outputs": pooler_output}
if new_token_ids or pooler_output is not None or kv_transfer_params:
# Add EngineCoreOutput for this Request.
outputs[request.client_index].append(
@@ -248,7 +264,7 @@ class DiffusionScheduler(OmniScheduler):
finish_reason=request.get_finished_reason(),
new_logprobs=new_logprobs,
new_prompt_logprobs_tensors=prompt_logprobs_tensors,
- pooling_output=pooler_output,
+ pooling_output=pooling_output_dict,
stop_reason=request.stop_reason,
events=request.take_events(),
kv_transfer_params=kv_transfer_params,
diff --git a/vllm_omni/core/sched/generation_scheduler.py b/vllm_omni/core/sched/generation_scheduler.py
new file mode 100644
index 000000000..7916c36ed
--- /dev/null
+++ b/vllm_omni/core/sched/generation_scheduler.py
@@ -0,0 +1,313 @@
+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 GenerationScheduler(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 "generation_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( # type: ignore[union-attr] # noqa: E501
+ req_id, new_token_ids
+ )
+
+ # 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)
+ # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput
+ pooling_output_dict = None
+ if pooler_output is not None:
+ # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type
+ # which expects Optional[dict[str, torch.Tensor]]
+ pooling_output_dict = {"model_outputs": pooler_output}
+ 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=pooling_output_dict,
+ 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
diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py
index dbfc20019..0d564fcda 100644
--- a/vllm_omni/core/sched/output.py
+++ b/vllm_omni/core/sched/output.py
@@ -9,6 +9,17 @@ from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
@dataclass
class OmniNewRequestData(NewRequestData):
+ """New request data for omni models with embeddings support.
+
+ Extends NewRequestData to include prompt embeddings and additional
+ information for direct transfer between pipeline stages.
+
+ Args:
+ prompt_embeds: Optional serialized prompt embeddings payload
+ additional_information: Optional serialized additional information
+ dictionary containing tensors or lists
+ """
+
# Optional serialized prompt embeddings
prompt_embeds: Optional[PromptEmbedsPayload] = None
# Optional serialized additional information
@@ -20,6 +31,15 @@ class OmniNewRequestData(NewRequestData):
request: Request,
block_ids: tuple[list[int], ...],
) -> "OmniNewRequestData":
+ """Create OmniNewRequestData from a Request object.
+
+ Args:
+ request: Request object to convert
+ block_ids: Tuple of block ID lists for KV cache allocation
+
+ Returns:
+ OmniNewRequestData instance with data from the request
+ """
return cls(
req_id=request.request_id,
prompt_token_ids=request.prompt_token_ids,
diff --git a/vllm_omni/core/sched/scheduler.py b/vllm_omni/core/sched/scheduler.py
index 3a30543c7..8d1fe8520 100644
--- a/vllm_omni/core/sched/scheduler.py
+++ b/vllm_omni/core/sched/scheduler.py
@@ -7,11 +7,11 @@ from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
class OmniScheduler(VLLMScheduler):
"""
- OmniScheduler: Scheduler for vLLM-omni multimodal processing.
+ 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.
+ specific to vLLM-Omni.
"""
# Ensure scheduled_new_reqs carry omni-specific payloads
diff --git a/vllm_omni/diffusion/__init__.py b/vllm_omni/diffusion/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
new file mode 100644
index 000000000..1b612264b
--- /dev/null
+++ b/vllm_omni/diffusion/data.py
@@ -0,0 +1,210 @@
+# adapted from sglang and fastvideo
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import enum
+import random
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.utils.network_utils import is_port_available
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class OmniDiffusionConfig:
+ # Model and path configuration (for convenience)
+ model: str
+
+ model_class_name: str | None = None
+
+ # Attention
+ # attention_backend: str = None
+
+ # Running mode
+ # mode: ExecutionMode = ExecutionMode.INFERENCE
+
+ # Workload type
+ # workload_type: WorkloadType = WorkloadType.T2V
+
+ # Cache strategy
+ cache_strategy: str = "none"
+
+ # Distributed executor backend
+ distributed_executor_backend: str = "mp"
+ nccl_port: int | None = None
+
+ # HuggingFace specific parameters
+ trust_remote_code: bool = False
+ revision: str | None = None
+
+ # Parallelism
+ num_gpus: int = 1
+ tp_size: int = -1
+ sp_degree: int = -1
+ # sequence parallelism
+ ulysses_degree: int | None = None
+ ring_degree: int | None = None
+ # data parallelism
+ # number of data parallelism groups
+ dp_size: int = 1
+ # number of gpu in a dp group
+ dp_degree: int = 1
+ # cfg parallel
+ enable_cfg_parallel: bool = False
+
+ hsdp_replicate_dim: int = 1
+ hsdp_shard_dim: int = -1
+ dist_timeout: int | None = None # timeout for torch.distributed
+
+ # pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False)
+
+ # LoRA parameters
+ # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated.
+ lora_path: str | None = None
+ lora_nickname: str = "default" # for swapping adapters in the pipeline
+ # can restrict layers to adapt, e.g. ["q_proj"]
+ # Will adapt only q, k, v, o by default.
+ lora_target_modules: list[str] | None = None
+
+ output_type: str = "pil"
+
+ # CPU offload parameters
+ dit_cpu_offload: bool = True
+ use_fsdp_inference: bool = False
+ text_encoder_cpu_offload: bool = True
+ image_encoder_cpu_offload: bool = True
+ vae_cpu_offload: bool = True
+ pin_cpu_memory: bool = True
+
+ # STA (Sliding Tile Attention) parameters
+ mask_strategy_file_path: str | None = None
+ # STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
+ skip_time_steps: int = 15
+
+ # Compilation
+ enable_torch_compile: bool = False
+
+ disable_autocast: bool = False
+
+ # VSA parameters
+ VSA_sparsity: float = 0.0 # inference/validation sparsity
+
+ # V-MoBA parameters
+ moba_config_path: str | None = None
+ # moba_config: dict[str, Any] = field(default_factory=dict)
+
+ # Master port for distributed inference
+ # TODO: do not hard code
+ master_port: int | None = None
+
+ # http server endpoint config, would be ignored in local mode
+ host: str | None = None
+ port: int | None = None
+
+ scheduler_port: int = 5555
+
+ # Stage verification
+ enable_stage_verification: bool = True
+
+ # Prompt text file for batch processing
+ prompt_file_path: str | None = None
+
+ # model paths for correct deallocation
+ model_paths: dict[str, str] = field(default_factory=dict)
+ model_loaded: dict[str, bool] = field(
+ default_factory=lambda: {
+ "transformer": True,
+ "vae": True,
+ }
+ )
+ override_transformer_cls_name: str | None = None
+
+ # # DMD parameters
+ # dmd_denoising_steps: List[int] | None = field(default=None)
+
+ # MoE parameters used by Wan2.2
+ boundary_ratio: float | None = None
+
+ # Logging
+ log_level: str = "info"
+
+ def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> int:
+ """
+ Find an available port with retry logic.
+
+ Args:
+ port: Initial port to check
+ port_inc: Port increment for each attempt
+ max_attempts: Maximum number of attempts to find an available port
+
+ Returns:
+ An available port number
+
+ Raises:
+ RuntimeError: If no available port is found after max_attempts
+ """
+ attempts = 0
+ original_port = port
+
+ while attempts < max_attempts:
+ if is_port_available(port):
+ if attempts > 0:
+ logger.info(f"Port {original_port} was unavailable, using port {port} instead")
+ return port
+
+ attempts += 1
+ if port < 60000:
+ port += port_inc
+ else:
+ # Wrap around with randomization to avoid collision
+ port = 5000 + random.randint(0, 1000)
+
+ raise RuntimeError(
+ f"Failed to find available port after {max_attempts} attempts (started from port {original_port})"
+ )
+
+ def __post_init__(self):
+ # TODO: remove hard code
+ initial_master_port = (self.master_port or 30005) + random.randint(0, 100)
+ self.master_port = self.settle_port(initial_master_port, 37)
+
+ @classmethod
+ def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
+ return cls(**kwargs)
+
+
+@dataclass
+class DiffusionOutput:
+ """
+ Final output (after pipeline completion)
+ """
+
+ output: torch.Tensor | None = None
+ trajectory_timesteps: list[torch.Tensor] | None = None
+ trajectory_latents: torch.Tensor | None = None
+ trajectory_decoded: list[torch.Tensor] | None = None
+ error: str | None = None
+
+ post_process_func: Callable[..., Any] | None = None
+
+ # logged timings info, directly from Req.timings
+ # timings: Optional["RequestTimings"] = None
+
+
+class AttentionBackendEnum(enum.Enum):
+ FA = enum.auto()
+ SLIDING_TILE_ATTN = enum.auto()
+ TORCH_SDPA = enum.auto()
+ SAGE_ATTN = enum.auto()
+ SAGE_ATTN_THREE = enum.auto()
+ VIDEO_SPARSE_ATTN = enum.auto()
+ VMOBA_ATTN = enum.auto()
+ AITER = enum.auto()
+ NO_ATTENTION = enum.auto()
+
+ def __str__(self):
+ return self.name.lower()
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
new file mode 100644
index 000000000..cd1e20458
--- /dev/null
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import multiprocessing as mp
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.registry import get_diffusion_post_process_func
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.scheduler import scheduler
+from vllm_omni.diffusion.worker.gpu_worker import WorkerProc
+
+logger = init_logger(__name__)
+
+
+class DiffusionEngine:
+ """The diffusion engine for vLLM-Omni diffusion models."""
+
+ def __init__(self, od_config: OmniDiffusionConfig):
+ """Initialize the diffusion engine.
+
+ Args:
+ config: The configuration for the diffusion engine.
+ """
+ self.od_config = od_config
+
+ self.post_process_func = get_diffusion_post_process_func(od_config)
+
+ self._make_client()
+
+ def step(self, requests: list[OmniDiffusionRequest]):
+ try:
+ output = self.add_req_and_wait_for_response(requests)
+ if output.error:
+ raise Exception(f"{output.error}")
+ logger.info("Generation completed successfully.")
+ return self.post_process_func(output.output)
+ except Exception as e:
+ logger.error(f"Generation failed: {e}")
+ return None
+
+ @staticmethod
+ def make_engine(config: OmniDiffusionConfig) -> "DiffusionEngine":
+ """Factory method to create a DiffusionEngine instance.
+
+ Args:
+ config: The configuration for the diffusion engine.
+
+ Returns:
+ An instance of DiffusionEngine.
+ """
+ return DiffusionEngine(config)
+
+ def _make_client(self):
+ # TODO rename it
+ scheduler.initialize(self.od_config)
+
+ # Get the broadcast handle from the initialized scheduler
+ broadcast_handle = scheduler.get_broadcast_handle()
+
+ processes, result_handle = self._launch_workers(
+ broadcast_handle=broadcast_handle,
+ )
+
+ if result_handle is not None:
+ scheduler.initialize_result_queue(result_handle)
+ else:
+ logger.error("Failed to get result queue handle from workers")
+
+ def _launch_workers(self, broadcast_handle):
+ od_config = self.od_config
+ logger.info("Starting server...")
+
+ num_gpus = od_config.num_gpus
+ mp.set_start_method("spawn", force=True)
+ processes = []
+
+ # Launch all worker processes
+ scheduler_pipe_readers = []
+ scheduler_pipe_writers = []
+
+ for i in range(num_gpus):
+ reader, writer = mp.Pipe(duplex=False)
+ scheduler_pipe_writers.append(writer)
+ process = mp.Process(
+ target=WorkerProc.worker_main,
+ args=(
+ i, # rank
+ od_config,
+ writer,
+ broadcast_handle,
+ ),
+ name=f"DiffusionWorker-{i}",
+ daemon=True,
+ )
+ scheduler_pipe_readers.append(reader)
+ process.start()
+ processes.append(process)
+
+ # Wait for all workers to be ready
+ scheduler_infos = []
+ result_handle = None
+ for writer in scheduler_pipe_writers:
+ writer.close()
+
+ for i, reader in enumerate(scheduler_pipe_readers):
+ try:
+ data = reader.recv()
+ except EOFError:
+ logger.error(f"Rank {i} scheduler is dead. Please check if there are relevant logs.")
+ processes[i].join()
+ logger.error(f"Exit code: {processes[i].exitcode}")
+ raise
+
+ if data["status"] != "ready":
+ raise RuntimeError("Initialization failed. Please see the error messages above.")
+
+ if i == 0:
+ result_handle = data.get("result_handle")
+
+ scheduler_infos.append(data)
+ reader.close()
+
+ logger.debug("All workers are ready")
+
+ return processes, result_handle
+
+ def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
+ return scheduler.add_req(requests)
diff --git a/vllm_omni/diffusion/distributed/__init__.py b/vllm_omni/diffusion/distributed/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_omni/diffusion/distributed/utils.py b/vllm_omni/diffusion/distributed/utils.py
new file mode 100644
index 000000000..fb930f21a
--- /dev/null
+++ b/vllm_omni/diffusion/distributed/utils.py
@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+import torch
+
+
+def get_local_device() -> torch.device:
+ """Return the torch device for the current rank."""
+ return torch.device(f"cuda:{os.environ.get('LOCAL_RANK', 0)}")
diff --git a/vllm_omni/diffusion/layers/__init__.py b/vllm_omni/diffusion/layers/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
new file mode 100644
index 000000000..cbedc5c31
--- /dev/null
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
@@ -0,0 +1,777 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import inspect
+import json
+import logging
+import math
+import os
+from typing import Any, Optional, Union
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.autoencoders.autoencoder_kl_qwenimage import (
+ AutoencoderKLQwenImage,
+)
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
+ FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer
+from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
+ QwenImageTransformer2DModel,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.model_executor.model_loader.weight_utils import (
+ download_weights_from_hf_specific,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_qwen_image_post_process_func(
+ od_config: OmniDiffusionConfig,
+):
+ model_name = od_config.model
+ if os.path.exists(model_name):
+ model_path = model_name
+ else:
+ model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+ vae_config_path = os.path.join(model_path, "vae/config.json")
+ with open(vae_config_path) as f:
+ vae_config = json.load(f)
+ vae_scale_factor = 2 ** len(vae_config["temporal_downsample"]) if "temporal_downsample" in vae_config else 8
+
+ image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+
+ def post_process_func(
+ images: torch.Tensor,
+ ):
+ return image_processor.postprocess(images)
+
+ return post_process_func
+
+
+def calculate_shift(
+ image_seq_len,
+ base_seq_len: int = 256,
+ max_seq_len: int = 4096,
+ base_shift: float = 0.5,
+ max_shift: float = 1.15,
+):
+ m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+ b = base_shift - m * base_seq_len
+ mu = image_seq_len * m + b
+ return mu
+
+
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: Optional[int] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ timesteps: Optional[list[int]] = None,
+ sigmas: Optional[list[float]] = None,
+ **kwargs,
+) -> tuple[torch.Tensor, int]:
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+def get_timestep_embedding(
+ timesteps: torch.Tensor,
+ embedding_dim: int,
+ flip_sin_to_cos: bool = False,
+ downscale_freq_shift: float = 1,
+ scale: float = 1,
+ max_period: int = 10000,
+) -> torch.Tensor:
+ """
+ This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+
+ Args
+ timesteps (torch.Tensor):
+ a 1-D Tensor of N indices, one per batch element. These may be fractional.
+ embedding_dim (int):
+ the dimension of the output.
+ flip_sin_to_cos (bool):
+ Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
+ downscale_freq_shift (float):
+ Controls the delta between frequencies between dimensions
+ scale (float):
+ Scaling factor applied to the embeddings.
+ max_period (int):
+ Controls the maximum frequency of the embeddings
+ Returns
+ torch.Tensor: an [N x dim] Tensor of positional embeddings.
+ """
+ assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+ half_dim = embedding_dim // 2
+ exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
+ exponent = exponent / (half_dim - downscale_freq_shift)
+
+ emb = torch.exp(exponent).to(timesteps.dtype)
+ emb = timesteps[:, None].float() * emb[None, :]
+
+ # scale embeddings
+ emb = scale * emb
+
+ # concat sine and cosine embeddings
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+ # flip sine and cosine embeddings
+ if flip_sin_to_cos:
+ emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+ # zero pad
+ if embedding_dim % 2 == 1:
+ emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+def apply_rotary_emb_qwen(
+ x: torch.Tensor,
+ freqs_cis: Union[torch.Tensor, tuple[torch.Tensor]],
+ use_real: bool = True,
+ use_real_unbind_dim: int = -1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+ to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+ reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+ tensors contain rotary embeddings and are returned as real tensors.
+
+ Args:
+ x (`torch.Tensor`):
+ Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply
+ freqs_cis (`tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+ Returns:
+ tuple[torch.Tensor, torch.Tensor]: tuple of modified query tensor and key tensor with rotary embeddings.
+ """
+ if use_real:
+ cos, sin = freqs_cis # [S, D]
+ cos = cos[None, None]
+ sin = sin[None, None]
+ cos, sin = cos.to(x.device), sin.to(x.device)
+
+ if use_real_unbind_dim == -1:
+ # Used for flux, cogvideox, hunyuan-dit
+ x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
+ x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+ elif use_real_unbind_dim == -2:
+ # Used for Stable Audio, OmniGen, CogView4 and Cosmos
+ x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, S, H, D//2]
+ x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+ else:
+ raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+ out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+ return out
+ else:
+ x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+ freqs_cis = freqs_cis.unsqueeze(1)
+ x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+ return x_out.type_as(x)
+
+
+class QwenImagePipeline(
+ nn.Module,
+):
+ def __init__(
+ self,
+ *,
+ od_config: OmniDiffusionConfig = None,
+ prefix: str = "",
+ ):
+ super().__init__()
+ self.device = get_local_device()
+ model = od_config.model
+ self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model, subfolder="scheduler")
+ logger.info("Loaded Qwen-Image scheduler successfully")
+
+ with set_default_torch_dtype(torch.bfloat16):
+ self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(model, subfolder="text_encoder")
+ logger.info("Loaded Qwen-Image text encoder successfully")
+ self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae").to(self.device)
+ logger.info("Loaded Qwen-Image VAE successfully")
+ self.transformer = QwenImageTransformer2DModel()
+ logger.info("Initialized Qwen-Image transformer successfully.")
+ self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer")
+ logger.info("Loaded Qwen-Image tokenizer successfully.")
+
+ self.stage = None
+
+ self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
+ # QwenImage latents are turned into 2x2 patches and packed.
+ # This means the latent width and height has to be divisible
+ # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+ # self.image_processor = VaeImageProcessor(
+ # vae_scale_factor=self.vae_scale_factor * 2
+ # )
+ self.tokenizer_max_length = 1024
+ self.prompt_template_encode = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" # noqa: E501
+ self.prompt_template_encode_start_idx = 34
+ self.default_sample_size = 128
+
+ def check_inputs(
+ self,
+ prompt,
+ height,
+ width,
+ negative_prompt=None,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ prompt_embeds_mask=None,
+ negative_prompt_embeds_mask=None,
+ callback_on_step_end_tensor_inputs=None,
+ max_sequence_length=None,
+ ):
+ if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
+ logger.warning(
+ f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} "
+ "but are {height} and {width}. Dimensions will be resized accordingly"
+ )
+
+ # if callback_on_step_end_tensor_inputs is not None and not all(
+ # k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+ # ):
+ # raise ValueError(
+ # f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs},
+ # but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+ # )
+
+ if prompt is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt is None and prompt_embeds is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+ )
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if negative_prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if prompt_embeds is not None and prompt_embeds_mask is None:
+ raise ValueError(
+ "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. "
+ "Make sure to generate `prompt_embeds_mask` from the same text encoder "
+ "that was used to generate `prompt_embeds`."
+ )
+ if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
+ raise ValueError(
+ "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. "
+ "Make sure to generate `negative_prompt_embeds_mask` from the same text encoder "
+ "that was used to generate `negative_prompt_embeds`."
+ )
+
+ if max_sequence_length is not None and max_sequence_length > 1024:
+ raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+
+ def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
+ bool_mask = mask.bool()
+ valid_lengths = bool_mask.sum(dim=1)
+ selected = hidden_states[bool_mask]
+ split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
+
+ return split_result
+
+ def _get_qwen_prompt_embeds(
+ self,
+ prompt: Union[str, list[str]] = None,
+ dtype: Optional[torch.dtype] = None,
+ ):
+ dtype = dtype or self.text_encoder.dtype
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ template = self.prompt_template_encode
+ drop_idx = self.prompt_template_encode_start_idx
+ txt = [template.format(e) for e in prompt]
+ txt_tokens = self.tokenizer(
+ txt,
+ max_length=self.tokenizer_max_length + drop_idx,
+ padding=True,
+ truncation=True,
+ return_tensors="pt",
+ ).to(self.device)
+ # print(f"attention mask: {txt_tokens.attention_mask}")
+ encoder_hidden_states = self.text_encoder(
+ input_ids=txt_tokens.input_ids,
+ attention_mask=txt_tokens.attention_mask,
+ output_hidden_states=True,
+ )
+ hidden_states = encoder_hidden_states.hidden_states[-1]
+ split_hidden_states = self._extract_masked_hidden(hidden_states, txt_tokens.attention_mask)
+ split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+ attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+ max_seq_len = max([e.size(0) for e in split_hidden_states])
+ prompt_embeds = torch.stack(
+ [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+ )
+ encoder_attention_mask = torch.stack(
+ [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+ )
+
+ prompt_embeds = prompt_embeds.to(dtype=dtype)
+
+ return prompt_embeds, encoder_attention_mask
+
+ def encode_prompt(
+ self,
+ prompt: Union[str, list[str]],
+ num_images_per_prompt: int = 1,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds_mask: Optional[torch.Tensor] = None,
+ max_sequence_length: int = 1024,
+ ):
+ r"""
+
+ Args:
+ prompt (`str` or `list[str]`, *optional*):
+ prompt to be encoded
+ device: (`torch.device`):
+ torch device
+ num_images_per_prompt (`int`):
+ number of images that should be generated per prompt
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ """
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
+
+ if prompt_embeds is None:
+ prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt)
+
+ prompt_embeds = prompt_embeds[:, :max_sequence_length]
+ prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
+
+ _, seq_len, _ = prompt_embeds.shape
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+ prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
+ prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+ return prompt_embeds, prompt_embeds_mask
+
+ @staticmethod
+ def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+ latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+ latents = latents.permute(0, 2, 4, 1, 3, 5)
+ latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+ return latents
+
+ @staticmethod
+ def _unpack_latents(latents, height, width, vae_scale_factor):
+ batch_size, num_patches, channels = latents.shape
+
+ # VAE applies 8x compression on images but we must also account for packing which requires
+ # latent height and width to be divisible by 2.
+ height = 2 * (int(height) // (vae_scale_factor * 2))
+ width = 2 * (int(width) // (vae_scale_factor * 2))
+
+ latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+ latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+ latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width)
+
+ return latents
+
+ def prepare_latents(
+ self,
+ batch_size,
+ num_channels_latents,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ latents=None,
+ ) -> torch.Tensor:
+ # generator=torch.Generator(device="cuda").manual_seed(42)
+ # VAE applies 8x compression on images but we must also account for packing which requires
+ # latent height and width to be divisible by 2.
+ height = 2 * (int(height) // (self.vae_scale_factor * 2))
+ width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+ shape = (batch_size, 1, num_channels_latents, height, width)
+
+ if latents is not None:
+ return latents.to(device=device, dtype=dtype)
+
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+
+ return latents
+
+ def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len):
+ sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+ # image_seq_len = latents.shape[1]
+ mu = calculate_shift(
+ image_seq_len,
+ self.scheduler.config.get("base_image_seq_len", 256),
+ self.scheduler.config.get("max_image_seq_len", 4096),
+ self.scheduler.config.get("base_shift", 0.5),
+ self.scheduler.config.get("max_shift", 1.15),
+ )
+ timesteps, num_inference_steps = retrieve_timesteps(
+ self.scheduler,
+ num_inference_steps,
+ sigmas=sigmas,
+ mu=mu,
+ )
+ return timesteps, num_inference_steps
+
+ @property
+ def guidance_scale(self):
+ return self._guidance_scale
+
+ @property
+ def attention_kwargs(self):
+ return self._attention_kwargs
+
+ @property
+ def num_timesteps(self):
+ return self._num_timesteps
+
+ @property
+ def current_timestep(self):
+ return self._current_timestep
+
+ @property
+ def interrupt(self):
+ return self._interrupt
+
+ def diffuse(
+ self,
+ prompt_embeds,
+ prompt_embeds_mask,
+ negative_prompt_embeds,
+ negative_prompt_embeds_mask,
+ latents,
+ img_shapes,
+ txt_seq_lens,
+ negative_txt_seq_lens,
+ timesteps,
+ do_true_cfg,
+ guidance,
+ true_cfg_scale,
+ ):
+ self.scheduler.set_begin_index(0)
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+ self._current_timestep = t
+ # broadcast to batch dimension and place on same device/dtype as latents
+ timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
+ noise_pred = self.transformer(
+ hidden_states=latents,
+ timestep=timestep / 1000,
+ guidance=guidance,
+ encoder_hidden_states_mask=prompt_embeds_mask,
+ encoder_hidden_states=prompt_embeds,
+ img_shapes=img_shapes,
+ txt_seq_lens=txt_seq_lens,
+ attention_kwargs=self.attention_kwargs,
+ return_dict=False,
+ )[0]
+ if do_true_cfg:
+ neg_noise_pred = self.transformer(
+ hidden_states=latents,
+ timestep=timestep / 1000,
+ guidance=guidance,
+ encoder_hidden_states_mask=negative_prompt_embeds_mask,
+ encoder_hidden_states=negative_prompt_embeds,
+ img_shapes=img_shapes,
+ txt_seq_lens=negative_txt_seq_lens,
+ attention_kwargs=self.attention_kwargs,
+ return_dict=False,
+ )[0]
+ comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+ cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+ noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+ noise_pred = comb_pred * (cond_norm / noise_norm)
+ # compute the previous noisy sample x_t -> x_t-1
+ latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+ return latents
+
+ def forward(
+ self,
+ req: OmniDiffusionRequest,
+ prompt: Union[str, list[str]] = "",
+ negative_prompt: Union[str, list[str]] = "",
+ true_cfg_scale: float = 4.0,
+ height: int | None = None,
+ width: int | None = None,
+ num_inference_steps: int = 50,
+ sigmas: Optional[list[float]] = None,
+ guidance_scale: float = 1.0,
+ num_images_per_prompt: int = 1,
+ generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_embeds_mask: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
+ output_type: Optional[str] = "pil",
+ attention_kwargs: Optional[dict[str, Any]] = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ max_sequence_length: int = 512,
+ ) -> DiffusionOutput:
+ # # TODO: only support single prompt now
+ # if req.prompt is not None:
+ # prompt = req.prompt[0] if isinstance(req.prompt, list) else req.prompt
+ prompt = req.prompt if req.prompt is not None else prompt
+ negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
+ height = req.height or self.default_sample_size * self.vae_scale_factor
+ width = req.width or self.default_sample_size * self.vae_scale_factor
+ num_inference_steps = req.num_inference_steps or num_inference_steps
+ generator = req.generator or generator
+ true_cfg_scale = req.true_cfg_scale or true_cfg_scale
+ req_num_outputs = getattr(req, "num_outputs_per_prompt", None)
+ if req_num_outputs and req_num_outputs > 0:
+ num_images_per_prompt = req_num_outputs
+ # 1. check inputs
+ # 2. encode prompts
+ # 3. prepare latents and timesteps
+ # 4. diffusion process
+ # 5. decode latents
+ # 6. post-process outputs
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ negative_prompt,
+ prompt_embeds,
+ negative_prompt_embeds,
+ prompt_embeds_mask,
+ negative_prompt_embeds_mask,
+ callback_on_step_end_tensor_inputs,
+ max_sequence_length,
+ )
+
+ self._guidance_scale = guidance_scale
+ self._attention_kwargs = attention_kwargs
+ self._current_timestep = None
+ self._interrupt = False
+
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ has_neg_prompt = negative_prompt is not None or (
+ negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
+ )
+
+ do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+ prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+ prompt=prompt,
+ prompt_embeds=prompt_embeds,
+ prompt_embeds_mask=prompt_embeds_mask,
+ num_images_per_prompt=num_images_per_prompt,
+ max_sequence_length=max_sequence_length,
+ )
+ if do_true_cfg:
+ negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+ prompt=negative_prompt,
+ prompt_embeds=negative_prompt_embeds,
+ prompt_embeds_mask=negative_prompt_embeds_mask,
+ num_images_per_prompt=num_images_per_prompt,
+ max_sequence_length=max_sequence_length,
+ )
+
+ num_channels_latents = self.transformer.in_channels // 4
+ latents = self.prepare_latents(
+ batch_size * num_images_per_prompt,
+ num_channels_latents,
+ height,
+ width,
+ prompt_embeds.dtype,
+ self.device,
+ generator,
+ latents,
+ )
+ img_shapes = [
+ [
+ (
+ 1,
+ height // self.vae_scale_factor // 2,
+ width // self.vae_scale_factor // 2,
+ )
+ ]
+ ] * batch_size
+
+ timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
+ # num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+ self._num_timesteps = len(timesteps)
+
+ # handle guidance
+ if self.transformer.guidance_embeds:
+ guidance = torch.full([1], guidance_scale, dtype=torch.float32)
+ guidance = guidance.expand(latents.shape[0])
+ else:
+ guidance = None
+
+ if self.attention_kwargs is None:
+ self._attention_kwargs = {}
+
+ txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
+ negative_txt_seq_lens = (
+ negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
+ )
+ # print inputp params
+
+ latents = self.diffuse(
+ prompt_embeds,
+ prompt_embeds_mask,
+ negative_prompt_embeds,
+ negative_prompt_embeds_mask,
+ latents,
+ img_shapes,
+ txt_seq_lens,
+ negative_txt_seq_lens,
+ timesteps,
+ do_true_cfg,
+ guidance,
+ true_cfg_scale,
+ )
+
+ self._current_timestep = None
+ if output_type == "latent":
+ image = latents
+ else:
+ latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+ latents = latents.to(self.vae.dtype)
+ latents_mean = (
+ torch.tensor(self.vae.config.latents_mean)
+ .view(1, self.vae.config.z_dim, 1, 1, 1)
+ .to(latents.device, latents.dtype)
+ )
+ latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+ latents.device, latents.dtype
+ )
+ latents = latents / latents_std + latents_mean
+ image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+ # processed_image = self.image_processor.postprocess(image, output_type=output_type)
+
+ return DiffusionOutput(output=image)
+
+ def load_weights(self):
+ self.load_transformer()
+
+ # TODO: find a better way to load
+ def load_transformer(self):
+ import glob
+ import os
+
+ # Define the weight iterator
+ def weight_iterator(transformer_path):
+ if not os.path.exists(transformer_path):
+ logger.warning(f"Path {transformer_path} does not exist.")
+ return
+
+ # Look for safetensors first
+ safetensors_files = glob.glob(os.path.join(transformer_path, "*.safetensors"))
+ if safetensors_files:
+ try:
+ from safetensors.torch import load_file
+ except ImportError:
+ logger.warning("safetensors not installed, cannot load .safetensors files.")
+ return
+
+ for file_path in safetensors_files:
+ state_dict = load_file(file_path)
+ for name, tensor in state_dict.items():
+ yield name, tensor
+ else:
+ # Fallback to bin
+ bin_files = glob.glob(os.path.join(transformer_path, "*.bin"))
+ for file_path in bin_files:
+ state_dict = torch.load(file_path)
+ for name, tensor in state_dict.items():
+ yield name, tensor
+
+ model_name = "Qwen/Qwen-Image"
+ try:
+ # Use the arguments as verified by the user
+ model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+ transformer_path = os.path.join(model_path, "transformer")
+ self.transformer.load_weights(weight_iterator(transformer_path))
+
+ except Exception as e:
+ print(f"An error occurred: {e}")
+ raise e
diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
new file mode 100644
index 000000000..c0ef5fd12
--- /dev/null
+++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
@@ -0,0 +1,700 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import functools
+from collections.abc import Iterable
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention import FeedForward
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+
+# TODO replace this with vLLM implementation
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.normalization import AdaLayerNormContinuous
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
+
+logger = init_logger(__name__)
+
+
+def apply_rotary_emb_qwen(
+ x: torch.Tensor,
+ freqs_cis: Union[torch.Tensor, tuple[torch.Tensor]],
+ use_real: bool = True,
+ use_real_unbind_dim: int = -1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+ to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+ reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+ tensors contain rotary embeddings and are returned as real tensors.
+
+ Args:
+ x (`torch.Tensor`):
+ Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply
+ freqs_cis (`tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+ Returns:
+ tuple[torch.Tensor, torch.Tensor]: tuple of modified query tensor and key tensor with rotary embeddings.
+ """
+ if use_real:
+ cos, sin = freqs_cis # [S, D]
+ cos = cos[None, None]
+ sin = sin[None, None]
+ cos, sin = cos.to(x.device), sin.to(x.device)
+
+ if use_real_unbind_dim == -1:
+ # Used for flux, cogvideox, hunyuan-dit
+ x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
+ x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+ elif use_real_unbind_dim == -2:
+ # Used for Stable Audio, OmniGen, CogView4 and Cosmos
+ x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, S, H, D//2]
+ x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+ else:
+ raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+ out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+ return out
+ else:
+ x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+ freqs_cis = freqs_cis.unsqueeze(1)
+ x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+ return x_out.type_as(x)
+
+
+class QwenTimestepProjEmbeddings(nn.Module):
+ def __init__(self, embedding_dim):
+ super().__init__()
+
+ self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000)
+ self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+ def forward(self, timestep, hidden_states):
+ timesteps_proj = self.time_proj(timestep)
+ timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype)) # (N, D)
+
+ conditioning = timesteps_emb
+
+ return conditioning
+
+
+class QwenEmbedRope(nn.Module):
+ def __init__(self, theta: int, axes_dim: list[int], scale_rope=False):
+ super().__init__()
+ self.theta = theta
+ self.axes_dim = axes_dim
+ pos_index = torch.arange(4096)
+ neg_index = torch.arange(4096).flip(0) * -1 - 1
+ self.pos_freqs = torch.cat(
+ [
+ self.rope_params(pos_index, self.axes_dim[0], self.theta),
+ self.rope_params(pos_index, self.axes_dim[1], self.theta),
+ self.rope_params(pos_index, self.axes_dim[2], self.theta),
+ ],
+ dim=1,
+ )
+ self.neg_freqs = torch.cat(
+ [
+ self.rope_params(neg_index, self.axes_dim[0], self.theta),
+ self.rope_params(neg_index, self.axes_dim[1], self.theta),
+ self.rope_params(neg_index, self.axes_dim[2], self.theta),
+ ],
+ dim=1,
+ )
+ self.rope_cache = {}
+
+ # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
+ self.scale_rope = scale_rope
+
+ def rope_params(self, index, dim, theta=10000):
+ """
+ Args:
+ index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
+ """
+ assert dim % 2 == 0
+ freqs = torch.outer(
+ index,
+ 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)),
+ )
+ freqs = torch.polar(torch.ones_like(freqs), freqs)
+ return freqs
+
+ def forward(self, video_fhw, txt_seq_lens, device):
+ """
+ Args: video_fhw: [frame, height, width] a list of 3 integers representing the shape of the video Args:
+ txt_length: [bs] a list of 1 integers representing the length of the text
+ """
+ if self.pos_freqs.device != device:
+ self.pos_freqs = self.pos_freqs.to(device)
+ self.neg_freqs = self.neg_freqs.to(device)
+
+ if isinstance(video_fhw, list):
+ video_fhw = video_fhw[0]
+ if not isinstance(video_fhw, list):
+ video_fhw = [video_fhw]
+
+ vid_freqs = []
+ max_vid_index = 0
+ for idx, fhw in enumerate(video_fhw):
+ frame, height, width = fhw
+ rope_key = f"{idx}_{height}_{width}"
+
+ if not torch.compiler.is_compiling():
+ if rope_key not in self.rope_cache:
+ self.rope_cache[rope_key] = self._compute_video_freqs(frame, height, width, idx)
+ video_freq = self.rope_cache[rope_key]
+ else:
+ video_freq = self._compute_video_freqs(frame, height, width, idx)
+ video_freq = video_freq.to(device)
+ vid_freqs.append(video_freq)
+
+ if self.scale_rope:
+ max_vid_index = max(height // 2, width // 2, max_vid_index)
+ else:
+ max_vid_index = max(height, width, max_vid_index)
+
+ max_len = max(txt_seq_lens)
+ txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...]
+ vid_freqs = torch.cat(vid_freqs, dim=0)
+
+ return vid_freqs, txt_freqs
+
+ @functools.cache
+ def _compute_video_freqs(self, frame, height, width, idx=0):
+ seq_lens = frame * height * width
+ freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+ freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+
+ freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1)
+ if self.scale_rope:
+ freqs_height = torch.cat(
+ [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]],
+ dim=0,
+ )
+ freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1)
+ freqs_width = torch.cat(
+ [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]],
+ dim=0,
+ )
+ freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1)
+ else:
+ freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1)
+ freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1)
+
+ freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
+ return freqs.clone().contiguous()
+
+
+class QwenImageCrossAttention(nn.Module):
+ def __init__(
+ self,
+ dim: int, # query_dim
+ num_heads: int,
+ head_dim: int,
+ window_size=(-1, -1),
+ added_kv_proj_dim: int = None,
+ out_bias: bool = True,
+ qk_norm=True, # rmsnorm
+ eps=1e-6,
+ pre_only=False,
+ context_pre_only: bool = False,
+ parallel_attention=False,
+ out_dim: int = None,
+ ) -> None:
+ assert dim % num_heads == 0
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.window_size = window_size
+ self.qk_norm = qk_norm
+ self.eps = eps
+ self.parallel_attention = parallel_attention
+
+ # layers
+ # self.to_q = ReplicatedLinear(dim, dim)
+ # self.to_k = ReplicatedLinear(dim, dim)
+ # self.to_v = ReplicatedLinear(dim, dim)
+ self.to_qkv = QKVParallelLinear(
+ hidden_size=dim,
+ head_size=self.head_dim,
+ total_num_heads=num_heads,
+ disable_tp=True,
+ )
+ self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+ self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads
+ self.inner_kv_dim = self.inner_dim
+ if added_kv_proj_dim is not None:
+ assert context_pre_only is not None
+ # self.add_k_proj = ReplicatedLinear(added_kv_proj_dim, self.inner_kv_dim, bias=True)
+ # self.add_v_proj = ReplicatedLinear(added_kv_proj_dim, self.inner_kv_dim, bias=True)
+ # self.add_q_proj = ReplicatedLinear(
+ # added_kv_proj_dim, self.inner_dim, bias=True
+ # )
+ self.add_kv_proj = QKVParallelLinear(
+ added_kv_proj_dim,
+ head_size=self.inner_kv_dim // self.num_heads,
+ total_num_heads=self.num_heads,
+ disable_tp=True,
+ )
+
+ if context_pre_only is not None and not context_pre_only:
+ self.to_add_out = ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias)
+ else:
+ self.to_add_out = None
+
+ if not pre_only:
+ self.to_out = nn.ModuleList([])
+ self.to_out.append(ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias))
+ else:
+ self.to_out = None
+
+ self.norm_added_q = RMSNorm(head_dim, eps=eps)
+ self.norm_added_k = RMSNorm(head_dim, eps=eps)
+
+ # self.attn = Attention(
+ # num_heads=num_heads,
+ # head_size=self.head_dim,
+ # scale=1.0, # softmax_scale
+ # attn_type=AttentionType.ENCODER,
+ # )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+ **cross_attention_kwargs,
+ ):
+ seq_len_txt = encoder_hidden_states.shape[1]
+
+ # Compute QKV for image stream (sample projections)
+ qkv, _ = self.to_qkv(hidden_states)
+ img_query, img_key, img_value = qkv.chunk(3, dim=-1)
+
+ # Compute QKV for text stream (context projections)
+ qkv, _ = self.add_kv_proj(encoder_hidden_states)
+ txt_query, txt_key, txt_value = qkv.chunk(3, dim=-1)
+
+ # Reshape for multi-head attention
+ img_query = img_query.unflatten(-1, (self.num_heads, -1))
+ img_key = img_key.unflatten(-1, (self.num_heads, -1))
+ img_value = img_value.unflatten(-1, (self.num_heads, -1))
+
+ txt_query = txt_query.unflatten(-1, (self.num_heads, -1))
+ txt_key = txt_key.unflatten(-1, (self.num_heads, -1))
+ txt_value = txt_value.unflatten(-1, (self.num_heads, -1))
+
+ # Apply QK normalization
+ img_query = self.norm_q(img_query)
+ img_key = self.norm_k(img_key)
+ txt_query = self.norm_added_q(txt_query)
+ txt_key = self.norm_added_k(txt_key)
+
+ # Apply RoPE
+ if image_rotary_emb is not None:
+ img_freqs, txt_freqs = image_rotary_emb
+ img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False)
+ img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False)
+ txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False)
+ txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False)
+
+ # Concatenate for joint attention
+ # Order: [text, image]
+ joint_query = torch.cat([txt_query, img_query], dim=1)
+ joint_key = torch.cat([txt_key, img_key], dim=1)
+ joint_value = torch.cat([txt_value, img_value], dim=1)
+
+ # Compute joint attention
+ # joint_hidden_states = self.attn(
+ # joint_query,
+ # joint_key,
+ # joint_value,
+ # )
+ joint_hidden_states = dispatch_attention_fn(
+ joint_query,
+ joint_key,
+ joint_value,
+ dropout_p=0.0,
+ is_causal=False,
+ )
+
+ # Reshape back
+ joint_hidden_states = joint_hidden_states.flatten(2, 3)
+ joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
+
+ # Split attention outputs back
+ txt_attn_output = joint_hidden_states[:, :seq_len_txt, :] # Text part
+ img_attn_output = joint_hidden_states[:, seq_len_txt:, :] # Image part
+
+ # Apply output projections
+ img_attn_output, _ = self.to_out[0](img_attn_output)
+ if len(self.to_out) > 1:
+ (img_attn_output,) = self.to_out[1](img_attn_output) # dropout
+
+ txt_attn_output, _ = self.to_add_out(txt_attn_output)
+
+ return img_attn_output, txt_attn_output
+
+
+class QwenImageTransformerBlock(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_attention_heads: int,
+ attention_head_dim: int,
+ qk_norm: str = "rms_norm",
+ eps: float = 1e-6,
+ ):
+ super().__init__()
+
+ self.dim = dim
+ self.num_attention_heads = num_attention_heads
+ self.attention_head_dim = attention_head_dim
+
+ # Image processing modules
+ self.img_mod = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(dim, 6 * dim, bias=True), # For scale, shift, gate for norm1 and norm2
+ )
+ self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ # self.attn = Attention(
+ # query_dim=dim,
+ # cross_attention_dim=None, # Enable cross attention for joint computation
+ # added_kv_proj_dim=dim, # Enable added KV projections for text stream
+ # dim_head=attention_head_dim,
+ # heads=num_attention_heads,
+ # out_dim=dim,
+ # context_pre_only=False,
+ # bias=True,
+ # processor=QwenDoubleStreamAttnProcessor2_0(),
+ # qk_norm=qk_norm,
+ # eps=eps,
+ # )
+ self.attn = QwenImageCrossAttention(
+ dim=dim,
+ num_heads=num_attention_heads,
+ added_kv_proj_dim=dim,
+ context_pre_only=False,
+ head_dim=attention_head_dim,
+ )
+ self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.img_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+ # Text processing modules
+ self.txt_mod = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(dim, 6 * dim, bias=True), # For scale, shift, gate for norm1 and norm2
+ )
+ self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ # Text doesn't need separate attention - it's handled by img_attn joint computation
+ self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.txt_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+ def _modulate(self, x, mod_params):
+ """Apply modulation to input tensor"""
+ shift, scale, gate = mod_params.chunk(3, dim=-1)
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1), gate.unsqueeze(1)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ encoder_hidden_states_mask: torch.Tensor,
+ temb: torch.Tensor,
+ image_rotary_emb: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ joint_attention_kwargs: Optional[dict[str, Any]] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ # Get modulation parameters for both streams
+ img_mod_params = self.img_mod(temb) # [B, 6*dim]
+ txt_mod_params = self.txt_mod(temb) # [B, 6*dim]
+
+ # Split modulation parameters for norm1 and norm2
+ img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1) # Each [B, 3*dim]
+ txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim]
+
+ # Process image stream - norm1 + modulation
+ img_normed = self.img_norm1(hidden_states)
+ img_modulated, img_gate1 = self._modulate(img_normed, img_mod1)
+
+ # Process text stream - norm1 + modulation
+ txt_normed = self.txt_norm1(encoder_hidden_states)
+ txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1)
+
+ # Use QwenAttnProcessor2_0 for joint attention computation
+ # This directly implements the DoubleStreamLayerMegatron logic:
+ # 1. Computes QKV for both streams
+ # 2. Applies QK normalization and RoPE
+ # 3. Concatenates and runs joint attention
+ # 4. Splits results back to separate streams
+ joint_attention_kwargs = joint_attention_kwargs or {}
+ attn_output = self.attn(
+ hidden_states=img_modulated, # Image stream (will be processed as "sample")
+ encoder_hidden_states=txt_modulated, # Text stream (will be processed as "context")
+ encoder_hidden_states_mask=encoder_hidden_states_mask,
+ image_rotary_emb=image_rotary_emb,
+ **joint_attention_kwargs,
+ )
+
+ # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
+ img_attn_output, txt_attn_output = attn_output
+
+ # Apply attention gates and add residual (like in Megatron)
+ hidden_states = hidden_states + img_gate1 * img_attn_output
+ encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
+
+ # Process image stream - norm2 + MLP
+ img_normed2 = self.img_norm2(hidden_states)
+ img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
+ img_mlp_output = self.img_mlp(img_modulated2)
+ hidden_states = hidden_states + img_gate2 * img_mlp_output
+
+ # Process text stream - norm2 + MLP
+ txt_normed2 = self.txt_norm2(encoder_hidden_states)
+ txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
+ txt_mlp_output = self.txt_mlp(txt_modulated2)
+ encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
+
+ # Clip to prevent overflow for fp16
+ if encoder_hidden_states.dtype == torch.float16:
+ encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+ if hidden_states.dtype == torch.float16:
+ hidden_states = hidden_states.clip(-65504, 65504)
+
+ return encoder_hidden_states, hidden_states
+
+
+class QwenImageTransformer2DModel(nn.Module):
+ """
+ The Transformer model introduced in Qwen.
+
+ Args:
+ patch_size (`int`, defaults to `2`):
+ Patch size to turn the input data into small patches.
+ in_channels (`int`, defaults to `64`):
+ The number of channels in the input.
+ out_channels (`int`, *optional*, defaults to `None`):
+ The number of channels in the output. If not specified, it defaults to `in_channels`.
+ num_layers (`int`, defaults to `60`):
+ The number of layers of dual stream DiT blocks to use.
+ attention_head_dim (`int`, defaults to `128`):
+ The number of dimensions to use for each attention head.
+ num_attention_heads (`int`, defaults to `24`):
+ The number of attention heads to use.
+ joint_attention_dim (`int`, defaults to `3584`):
+ The number of dimensions to use for the joint attention (embedding/channel dimension of
+ `encoder_hidden_states`).
+ guidance_embeds (`bool`, defaults to `False`):
+ Whether to use guidance embeddings for guidance-distilled variant of the model.
+ axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`):
+ The dimensions to use for the rotary positional embeddings.
+ """
+
+ # _supports_gradient_checkpointing = True
+ # _no_split_modules = ["QwenImageTransformerBlock"]
+ # _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+ # _repeated_blocks = ["QwenImageTransformerBlock"]
+ def __init__(
+ self,
+ patch_size: int = 2,
+ in_channels: int = 64,
+ out_channels: Optional[int] = 16,
+ num_layers: int = 60,
+ attention_head_dim: int = 128,
+ num_attention_heads: int = 24,
+ joint_attention_dim: int = 3584,
+ guidance_embeds: bool = False, # TODO: this should probably be removed
+ axes_dims_rope: tuple[int, int, int] = (16, 56, 56),
+ ):
+ super().__init__()
+ self.config = None
+ self.in_channels = in_channels
+ self.out_channels = out_channels or in_channels
+ self.inner_dim = num_attention_heads * attention_head_dim
+ self.guidance_embeds = guidance_embeds
+
+ self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True)
+
+ self.time_text_embed = QwenTimestepProjEmbeddings(embedding_dim=self.inner_dim)
+
+ self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
+
+ self.img_in = nn.Linear(in_channels, self.inner_dim)
+ self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim)
+
+ self.transformer_blocks = nn.ModuleList(
+ [
+ QwenImageTransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+ self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ encoder_hidden_states_mask: torch.Tensor = None,
+ timestep: torch.LongTensor = None,
+ img_shapes: Optional[list[tuple[int, int, int]]] = None,
+ txt_seq_lens: Optional[list[int]] = None,
+ guidance: torch.Tensor = None, # TODO: this should probably be removed
+ attention_kwargs: Optional[dict[str, Any]] = None,
+ return_dict: bool = True,
+ ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+ """
+ The [`QwenTransformer2DModel`] forward method.
+
+ Args:
+ hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+ Input `hidden_states`.
+ encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+ Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+ encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`):
+ Mask of the input conditions.
+ timestep ( `torch.LongTensor`):
+ Used to indicate denoising step.
+ attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+ `self.processor` in
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+ tuple.
+
+ Returns:
+ If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+ `tuple` where the first element is the sample tensor.
+ """
+ # if attention_kwargs is not None:
+ # attention_kwargs = attention_kwargs.copy()
+ # lora_scale = attention_kwargs.pop("scale", 1.0)
+ # else:
+ # lora_scale = 1.0
+
+ hidden_states = self.img_in(hidden_states)
+
+ # Ensure timestep tensor is on the same device and dtype as hidden_states
+ timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype)
+ encoder_hidden_states = self.txt_norm(encoder_hidden_states)
+ encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+ if guidance is not None:
+ guidance = guidance.to(hidden_states.dtype) * 1000
+
+ temb = (
+ self.time_text_embed(timestep, hidden_states)
+ if guidance is None
+ else self.time_text_embed(timestep, guidance, hidden_states)
+ )
+
+ image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
+
+ for index_block, block in enumerate(self.transformer_blocks):
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_hidden_states_mask=encoder_hidden_states_mask,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ joint_attention_kwargs=attention_kwargs,
+ )
+
+ # Use only the image part (hidden_states) from the dual-stream blocks
+ hidden_states = self.norm_out(hidden_states, temb)
+ output = self.proj_out(hidden_states)
+
+ return Transformer2DModelOutput(sample=output)
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+ # Collect incoming weights into a dict for flexible mapping and grouping
+ weights_map: dict[str, torch.Tensor] = {name: tensor for name, tensor in weights}
+
+ stacked_params_mapping = [
+ (
+ "to_qkv",
+ [
+ {"aliases": ["q_proj", "to_q", "q", "query"], "order": 0},
+ {"aliases": ["k_proj", "to_k", "k", "key"], "order": 1},
+ {"aliases": ["v_proj", "to_v", "v", "value"], "order": 2},
+ ],
+ ),
+ (
+ "add_kv_proj",
+ [
+ {"aliases": ["add_q_proj", "add_q"], "order": 0},
+ {"aliases": ["add_k_proj", "add_k"], "order": 1},
+ {"aliases": ["add_v_proj", "add_v"], "order": 2},
+ ],
+ ),
+ ]
+
+ # Helper that attempts to gather shards per base key using aliases and merges them
+ for target, shards in stacked_params_mapping:
+ base_candidates: dict[str, dict[str, dict[int, tuple[str, torch.Tensor]]]] = {}
+ for key in list(weights_map.keys()):
+ for shard in shards:
+ for alias in shard["aliases"]:
+ weight_suffix = f".{alias}.weight"
+ bias_suffix = f".{alias}.bias"
+ if key.endswith(weight_suffix):
+ base = key[: -len(weight_suffix)]
+ base_entry = base_candidates.setdefault(base, {"weight": {}, "bias": {}})
+ base_entry["weight"][shard["order"]] = (
+ key,
+ weights_map[key],
+ )
+ break
+ if key.endswith(bias_suffix):
+ base = key[: -len(bias_suffix)]
+ base_entry = base_candidates.setdefault(base, {"weight": {}, "bias": {}})
+ base_entry["bias"][shard["order"]] = (key, weights_map[key])
+ break
+
+ for base, entry in base_candidates.items():
+ if len(entry["weight"]) == len(shards):
+ ordered = [entry["weight"][order] for order in sorted(entry["weight"])]
+ combined = torch.cat([tensor for _, tensor in ordered], dim=0)
+ for key, _ in ordered:
+ weights_map.pop(key, None)
+ weights_map[base + f".{target}.weight"] = combined
+ if len(entry["bias"]) == len(shards):
+ ordered_bias = [entry["bias"][order] for order in sorted(entry["bias"])]
+ bcombined = torch.cat([tensor for _, tensor in ordered_bias], dim=0)
+ for key, _ in ordered_bias:
+ weights_map.pop(key, None)
+ weights_map[base + f".{target}.bias"] = bcombined
+
+ # Now perform the parameter copy using the possibly-updated weights_map
+ params_dict = dict(self.named_parameters())
+ for name, loaded_weight in weights_map.items():
+ # Handle potential name mismatches for FeedForward layers
+ if name.endswith(".proj.weight") and name not in params_dict:
+ mapped_name = name.replace(".proj.weight", ".weight")
+ elif name.endswith(".proj.bias") and name not in params_dict:
+ mapped_name = name.replace(".proj.bias", ".bias")
+ else:
+ mapped_name = name
+
+ if mapped_name in params_dict:
+ param = params_dict[mapped_name]
+ with torch.no_grad():
+ param.copy_(loaded_weight)
+ else:
+ logger.warning(f"Weight {name} (mapped to {mapped_name}) not found in model params.")
diff --git a/vllm_omni/diffusion/omni_diffusion.py b/vllm_omni/diffusion/omni_diffusion.py
new file mode 100644
index 000000000..c75d7da22
--- /dev/null
+++ b/vllm_omni/diffusion/omni_diffusion.py
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+from dataclasses import fields
+
+from vllm.logger import init_logger
+from vllm.transformers_utils.config import get_hf_file_to_dict
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+# TODO configure logging properly
+logging.basicConfig(level=logging.INFO)
+
+logger = init_logger(__name__)
+
+
+def prepare_requests(prompt: str | list[str], **kwargs):
+ field_names = {f.name for f in fields(OmniDiffusionRequest)}
+
+ init_kwargs = {"prompt": prompt}
+
+ for key, value in kwargs.items():
+ if key in field_names:
+ init_kwargs[key] = value
+
+ return OmniDiffusionRequest(**init_kwargs)
+
+
+class OmniDiffusion:
+ """
+ It is the main class to interact with vLLM-Omni diffusion models.
+ It acts as a high-level interface that prepares requests and
+ delegates the actual diffusion process to the DiffusionEngine.
+
+ You can pass either an `OmniDiffusionConfig` via `od_config`, or
+ pass kwargs such as `model="Qwen/Qwen-Image"`,
+ which will be forwarded to `OmniDiffusionConfig.from_kwargs`.
+ """
+
+ def __init__(self, od_config: OmniDiffusionConfig | None = None, **kwargs):
+ if od_config is None:
+ od_config = OmniDiffusionConfig.from_kwargs(**kwargs)
+ elif isinstance(od_config, dict):
+ od_config = OmniDiffusionConfig.from_kwargs(**od_config)
+
+ self.od_config = od_config
+
+ config_dict = get_hf_file_to_dict(
+ "model_index.json",
+ od_config.model,
+ )
+ od_config.model_class_name = config_dict.get("_class_name", None)
+
+ self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config)
+
+ def generate(
+ self,
+ prompt: str | list[str],
+ **kwargs,
+ ):
+ prompts = []
+ if isinstance(prompt, str):
+ prompts.append(prompt)
+ elif isinstance(prompt, list):
+ prompts.extend(prompt)
+ else:
+ raise ValueError("Prompt must be a string or a list of strings")
+
+ requests: list[OmniDiffusionRequest] = []
+ for p in prompts:
+ requests.append(
+ prepare_requests(
+ p,
+ **kwargs,
+ )
+ )
+ logger.info(f"Prepared {len(requests)} requests for generation.")
+ return self._run_engine(requests)
+
+ def _run_engine(self, requests: list[OmniDiffusionRequest]):
+ return self.engine.step(requests)
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
new file mode 100644
index 000000000..0d834b2a0
--- /dev/null
+++ b/vllm_omni/diffusion/registry.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+_DIFFUSION_MODELS = {
+ # arch:(mod_folder, mod_relname, cls_name)
+ "QwenImagePipeline": (
+ "qwen_image",
+ "pipeline_qwen_image",
+ "QwenImagePipeline",
+ ),
+}
+
+
+def initialize_model(
+ od_config: OmniDiffusionConfig,
+):
+ if od_config.model_class_name in _DIFFUSION_MODELS:
+ mod_folder, mod_relname, cls_name = _DIFFUSION_MODELS[od_config.model_class_name]
+ module_name = f"vllm_omni.diffusion.models.{mod_folder}.{mod_relname}"
+ module = importlib.import_module(module_name)
+ model_class = getattr(module, cls_name)
+ model = model_class(od_config=od_config, prefix=mod_relname)
+ return model
+ else:
+ raise ValueError(f"Model class {od_config.model_class_name} not found in diffusion model registry.")
+
+
+_DIFFUSION_POST_PROCESS_FUNCS = {
+ # arch: post_process_func
+ # `post_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
+ # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key
+ "QwenImagePipeline": "get_qwen_image_post_process_func",
+}
+
+
+def get_diffusion_post_process_func(od_config: OmniDiffusionConfig):
+ if od_config.model_class_name in _DIFFUSION_POST_PROCESS_FUNCS:
+ mod_folder, mod_relname, _ = _DIFFUSION_MODELS[od_config.model_class_name]
+ func_name = _DIFFUSION_POST_PROCESS_FUNCS[od_config.model_class_name]
+ module_name = f"vllm_omni.diffusion.models.{mod_folder}.{mod_relname}"
+ module = importlib.import_module(module_name)
+ post_process_func = getattr(module, func_name)
+ return post_process_func(od_config)
+ else:
+ raise ValueError(
+ f"Post process function for model class {od_config.model_class_name} not found in diffusion model registry."
+ )
diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py
new file mode 100644
index 000000000..100f8a50c
--- /dev/null
+++ b/vllm_omni/diffusion/request.py
@@ -0,0 +1,179 @@
+# adapted from sglang and fastvideo
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pprint
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import PIL.Image
+import torch
+
+
+@dataclass
+class OmniDiffusionRequest:
+ """
+ Complete state passed through the pipeline execution.
+
+ This dataclass contains all information needed during the diffusion pipeline
+ execution, allowing methods to update specific components without needing
+ to manage numerous individual parameters.
+ """
+
+ # TODO(will): double check that args are separate from server_args
+ # properly. Also maybe think about providing an abstraction for pipeline
+ # specific arguments.
+ # data_type: DataType
+
+ request_id: str | None = None
+
+ generator: torch.Generator | list[torch.Generator] | None = None
+
+ # Image inputs
+ image_path: str | None = None
+ # Image encoder hidden states
+ image_embeds: list[torch.Tensor] = field(default_factory=list)
+ pil_image: torch.Tensor | PIL.Image.Image | None = None
+ pixel_values: torch.Tensor | PIL.Image.Image | None = None
+ preprocessed_image: torch.Tensor | None = None
+
+ # Text inputs
+ prompt: str | list[str] | None = None
+ negative_prompt: str | list[str] | None = None
+ prompt_path: str | None = None
+ output_path: str = "outputs/"
+ # without extension
+ output_file_name: str | None = None
+ output_file_ext: str | None = None
+ # Primary encoder embeddings
+ prompt_embeds: list[torch.Tensor] | torch.Tensor = field(default_factory=list)
+ negative_prompt_embeds: list[torch.Tensor] | None = None
+ prompt_attention_mask: list[torch.Tensor] | None = None
+ negative_attention_mask: list[torch.Tensor] | None = None
+ clip_embedding_pos: list[torch.Tensor] | None = None
+ clip_embedding_neg: list[torch.Tensor] | None = None
+
+ pooled_embeds: list[torch.Tensor] = field(default_factory=list)
+ neg_pooled_embeds: list[torch.Tensor] = field(default_factory=list)
+
+ # Additional text-related parameters
+ max_sequence_length: int | None = None
+ prompt_template: dict[str, Any] | None = None
+ do_classifier_free_guidance: bool = False
+
+ # Batch info
+ num_outputs_per_prompt: int = 1
+ seed: int | None = None
+ seeds: list[int] | None = None
+
+ # Tracking if embeddings are already processed
+ is_prompt_processed: bool = False
+
+ # Latent tensors
+ latents: torch.Tensor | None = None
+ raw_latent_shape: torch.Tensor | None = None
+ noise_pred: torch.Tensor | None = None
+ image_latent: torch.Tensor | None = None
+
+ # Latent dimensions
+ height_latents: list[int] | int | None = None
+ width_latents: list[int] | int | None = None
+ num_frames: list[int] | int = 1 # Default for image models
+ num_frames_round_down: bool = False # Whether to round down num_frames if it's not divisible by num_gpus
+
+ # Original dimensions (before VAE scaling)
+ height: list[int] | int | None = None
+ width: list[int] | int | None = None
+ fps: list[int] | int | None = None
+ height_not_provided: bool = False
+ width_not_provided: bool = False
+
+ # Timesteps
+ timesteps: torch.Tensor | None = None
+ timestep: torch.Tensor | float | int | None = None
+ step_index: int | None = None
+ boundary_ratio: float | None = None
+
+ # Scheduler parameters
+ num_inference_steps: int = 50
+ guidance_scale: float = 1.0
+ guidance_scale_2: float | None = None
+ guidance_rescale: float = 0.0
+ eta: float = 0.0
+ sigmas: list[float] | None = None
+
+ true_cfg_scale: float | None = None # qwen-image specific now
+
+ n_tokens: int | None = None
+
+ # Other parameters that may be needed by specific schedulers
+ extra_step_kwargs: dict[str, Any] = field(default_factory=dict)
+
+ # Component modules (populated by the pipeline)
+ modules: dict[str, Any] = field(default_factory=dict)
+
+ return_trajectory_latents: bool = False
+ return_trajectory_decoded: bool = False
+ trajectory_timesteps: list[torch.Tensor] | None = None
+ trajectory_latents: torch.Tensor | None = None
+
+ # Extra parameters that might be needed by specific pipeline implementations
+ extra: dict[str, Any] = field(default_factory=dict)
+
+ # Misc
+ save_output: bool = True
+ return_frames: bool = False
+
+ # TeaCache parameters
+ enable_teacache: bool = False
+ # teacache_params: TeaCacheParams | WanTeaCacheParams | None = None
+
+ # STA parameters
+ STA_param: list | None = None
+ is_cfg_negative: bool = False
+ mask_search_final_result_pos: list[list] | None = None
+ mask_search_final_result_neg: list[list] | None = None
+
+ # VSA parameters
+ VSA_sparsity: float = 0.0
+ # perf_logger: PerformanceLogger | None = None
+
+ # stage logging
+ # logging_info: PipelineLoggingInfo = field(default_factory=PipelineLoggingInfo)
+
+ # profile
+ profile: bool = False
+ num_profiled_timesteps: int = 8
+
+ # debugging
+ debug: bool = False
+
+ # results
+ output: torch.Tensor | None = None
+
+ @property
+ def batch_size(self):
+ # Determine batch size
+ if isinstance(self.prompt, list):
+ batch_size = len(self.prompt)
+ elif self.prompt is not None:
+ batch_size = 1
+ else:
+ batch_size = self.prompt_embeds[0].shape[0]
+
+ # Adjust batch size for number of videos per prompt
+ batch_size *= self.num_outputs_per_prompt
+ return batch_size
+
+ def __post_init__(self):
+ """Initialize dependent fields after dataclass initialization."""
+ # Set do_classifier_free_guidance based on guidance scale and negative prompt
+ if self.guidance_scale > 1.0 and self.negative_prompt is not None:
+ self.do_classifier_free_guidance = True
+ if self.negative_prompt_embeds is None:
+ self.negative_prompt_embeds = []
+ if self.guidance_scale_2 is None:
+ self.guidance_scale_2 = self.guidance_scale
+
+ def __str__(self):
+ return pprint.pformat(asdict(self), indent=2, width=120)
diff --git a/vllm_omni/diffusion/scheduler.py b/vllm_omni/diffusion/scheduler.py
new file mode 100644
index 000000000..57c2c87df
--- /dev/null
+++ b/vllm_omni/diffusion/scheduler.py
@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import zmq
+from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+logger = init_logger(__name__)
+
+
+class Scheduler:
+ _instance = None
+
+ def __new__(cls, *args, **kwargs):
+ if not cls._instance:
+ cls._instance = super().__new__(cls)
+ return cls._instance
+
+ def initialize(self, od_config: OmniDiffusionConfig):
+ if hasattr(self, "context") and not self.context.closed:
+ logger.warning("SyncSchedulerClient is already initialized. Re-initializing.")
+ self.close()
+
+ self.od_config = od_config
+ self.context = zmq.Context() # Standard synchronous context
+
+ # Initialize MessageQueue for broadcasting requests
+ # Assuming all readers are local for now as per current launch_engine implementation
+ self.mq = MessageQueue(
+ n_reader=od_config.num_gpus,
+ n_local_reader=od_config.num_gpus,
+ local_reader_ranks=list(range(od_config.num_gpus)),
+ )
+
+ self.result_mq = None
+
+ def initialize_result_queue(self, handle):
+ # Initialize MessageQueue for receiving results
+ # We act as rank 0 reader for this queue
+ self.result_mq = MessageQueue.create_from_handle(handle, rank=0)
+ logger.info("SyncScheduler initialized result MessageQueue")
+
+ def get_broadcast_handle(self):
+ return self.mq.export_handle()
+
+ def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
+ """Sends a request to the scheduler and waits for the response."""
+ try:
+ # Broadcast request to all workers
+ self.mq.enqueue(requests)
+ # Wait for result from Rank 0 (or whoever sends it)
+ if self.result_mq is None:
+ raise RuntimeError("Result queue not initialized")
+
+ output = self.result_mq.dequeue()
+ return output
+ except zmq.error.Again:
+ logger.error("Timeout waiting for response from scheduler.")
+ raise TimeoutError("Scheduler did not respond in time.")
+
+ def close(self):
+ """Closes the socket and terminates the context."""
+ if hasattr(self, "context"):
+ self.context.term()
+
+
+# Singleton instance for easy access
+scheduler = Scheduler()
diff --git a/vllm_omni/diffusion/utils/__init__.py b/vllm_omni/diffusion/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_omni/diffusion/utils/hf_utils.py b/vllm_omni/diffusion/utils/hf_utils.py
new file mode 100644
index 000000000..61c7f1982
--- /dev/null
+++ b/vllm_omni/diffusion/utils/hf_utils.py
@@ -0,0 +1,20 @@
+from functools import lru_cache
+
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def load_diffusers_config(model_name) -> dict:
+ config = DiffusionPipeline.load_config(model_name)
+ return config
+
+
+@lru_cache
+def is_diffusion_model(model_name: str) -> bool:
+ try:
+ load_diffusers_config(model_name)
+ return True
+ except Exception:
+ return False
diff --git a/vllm_omni/diffusion/utils/network_utils.py b/vllm_omni/diffusion/utils/network_utils.py
new file mode 100644
index 000000000..c3085e8f9
--- /dev/null
+++ b/vllm_omni/diffusion/utils/network_utils.py
@@ -0,0 +1,19 @@
+import socket
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def is_port_available(port):
+ """Return whether a port is available."""
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+ try:
+ s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+ s.bind(("", port))
+ s.listen(1)
+ return True
+ except OSError:
+ return False
+ except OverflowError:
+ return False
diff --git a/vllm_omni/diffusion/worker/gpu_worker.py b/vllm_omni/diffusion/worker/gpu_worker.py
new file mode 100644
index 000000000..84ede7ffd
--- /dev/null
+++ b/vllm_omni/diffusion/worker/gpu_worker.py
@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import multiprocessing as mp
+import os
+
+import torch
+import zmq
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
+from vllm.distributed.parallel_state import (
+ init_distributed_environment,
+ initialize_model_parallel,
+)
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.registry import initialize_model
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+logger = init_logger(__name__)
+
+
+class GPUWorker:
+ """
+ A worker that executes the model on a single GPU.
+ """
+
+ def __init__(
+ self,
+ local_rank: int,
+ rank: int,
+ od_config: OmniDiffusionConfig,
+ ):
+ self.local_rank = local_rank
+ self.rank = rank
+ self.od_config = od_config
+ self.pipeline = None
+
+ self.init_device_and_model()
+
+ def init_device_and_model(self) -> None:
+ """Initialize the device and load the model."""
+ world_size = self.od_config.num_gpus
+ rank = self.rank
+ # Set environment variables for distributed initialization
+ os.environ["MASTER_ADDR"] = "localhost"
+ os.environ["MASTER_PORT"] = str(self.od_config.master_port)
+ os.environ["LOCAL_RANK"] = str(self.local_rank)
+ os.environ["RANK"] = str(rank)
+ os.environ["WORLD_SIZE"] = str(world_size)
+
+ device = torch.device(f"cuda:{rank}")
+ torch.cuda.set_device(device)
+
+ # hack
+ vllm_config = VllmConfig()
+ vllm_config.parallel_config.tensor_parallel_size = self.od_config.num_gpus
+ set_current_vllm_config(vllm_config)
+
+ init_distributed_environment(world_size=world_size, rank=rank)
+ initialize_model_parallel(tensor_model_parallel_size=world_size)
+
+ with device:
+ self.pipeline = initialize_model(self.od_config)
+ self.pipeline.load_weights()
+ self.pipeline.eval()
+ logger.info(f"Worker {self.rank}: Initialized device, model, and distributed environment.")
+ logger.info(f"Worker {self.rank}: Model loaded successfully.")
+
+ @torch.inference_mode()
+ def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
+ """
+ Execute a forward pass.
+ """
+ assert self.pipeline is not None
+ # TODO: dealing with first req for now
+ req = reqs[0]
+ output = self.pipeline.forward(req)
+ return output
+
+
+class WorkerProc:
+ """Wrapper that runs one Worker in a separate process."""
+
+ def __init__(
+ self,
+ od_config: OmniDiffusionConfig,
+ gpu_id: int,
+ broadcast_handle,
+ ):
+ self.od_config = od_config
+
+ # Inter-process Communication
+ self.context = zmq.Context(io_threads=2)
+
+ # Initialize MessageQueue reader from handle
+ self.mq = MessageQueue.create_from_handle(broadcast_handle, gpu_id)
+
+ self.result_mq = None
+ self.result_mq_handle = None
+
+ # Setup result sender (only for rank 0 for now, or whoever needs to reply)
+ # Assuming only rank 0 replies to scheduler as per original logic
+ if gpu_id == 0:
+ # Create MessageQueue for results (1 writer -> 1 reader)
+ # We assume the reader (SyncScheduler) will act as rank 0
+ self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
+ self.result_mq_handle = self.result_mq.export_handle()
+ logger.info(f"Worker {gpu_id} created result MessageQueue")
+
+ assert od_config.master_port is not None
+ worker = GPUWorker(
+ local_rank=gpu_id,
+ rank=gpu_id,
+ od_config=od_config,
+ )
+ self.worker = worker
+ self.gpu_id = gpu_id
+ self._running = True
+
+ def return_result(self, output: DiffusionOutput):
+ """
+ replies to client, only on rank 0
+ """
+ if self.result_mq is not None:
+ self.result_mq.enqueue(output)
+
+ def recv_reqs(self):
+ """
+ Receive requests from broadcast queue
+ """
+ return self.mq.dequeue()
+
+ # TODO: queueing, cancellation
+ def worker_busy_loop(self) -> None:
+ """Main busy loop for Multiprocessing Workers"""
+
+ logger.info(f"Worker {self.gpu_id} ready to receive requests via shared memory")
+
+ while self._running:
+ reqs = None
+ # 1: receive requests
+ try:
+ reqs = self.recv_reqs()
+ except Exception as e:
+ logger.error(
+ f"Error receiving requests in scheduler event loop: {e}",
+ exc_info=True,
+ )
+ continue
+
+ # 2: execute, make sure a reply is always sent
+ try:
+ output = self.worker.execute_model(reqs, self.od_config)
+ except Exception as e:
+ logger.error(
+ f"Error executing forward in event loop: {e}",
+ exc_info=True,
+ )
+ output = DiffusionOutput(error=str(e))
+
+ try:
+ self.return_result(output)
+ except zmq.ZMQError as e:
+ # Reply failed; log and keep loop alive to accept future requests
+ logger.error(f"ZMQ error sending reply: {e}")
+ continue
+
+ logger.info("event loop terminated.")
+ # if self.result_sender is not None:
+ # self.result_sender.close()
+ self.context.term()
+
+ @staticmethod
+ def worker_main(
+ rank: int,
+ od_config: OmniDiffusionConfig,
+ pipe_writer: mp.connection.Connection,
+ broadcast_handle,
+ ) -> None:
+ """Worker initialization and execution loops."""
+
+ worker_proc = WorkerProc(
+ od_config,
+ gpu_id=rank,
+ broadcast_handle=broadcast_handle,
+ )
+ logger.info(f"Worker {rank}: Scheduler loop started.")
+ pipe_writer.send(
+ {
+ "status": "ready",
+ "result_handle": worker_proc.result_mq_handle if rank == 0 else None,
+ }
+ )
+ worker_proc.worker_busy_loop()
+ logger.info(f"Worker {rank}: Shutdown complete.")
diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py
index 2de831388..8ce5f043a 100644
--- a/vllm_omni/engine/__init__.py
+++ b/vllm_omni/engine/__init__.py
@@ -1,11 +1,22 @@
"""
-Engine components for vLLM-omni.
+Engine components for vLLM-Omni.
"""
-from typing import Any, Optional
+import time
+from collections.abc import Mapping
+from typing import Any, Optional, Union
import msgspec
-from vllm.v1.engine import EngineCoreRequest
+import torch
+from vllm.v1.engine import (
+ EngineCoreEvent,
+ EngineCoreRequest,
+ FinishReason,
+ LogprobsLists,
+ LogprobsTensors,
+ SchedulerStats,
+ UtilityOutput,
+)
class PromptEmbedsPayload(msgspec.Struct):
@@ -49,7 +60,79 @@ class AdditionalInformationPayload(msgspec.Struct):
class OmniEngineCoreRequest(EngineCoreRequest):
+ """Engine core request for omni models with embeddings support.
+
+ Extends the base EngineCoreRequest with support for prompt embeddings
+ and additional information payloads, enabling direct transfer of
+ pre-computed embeddings between pipeline stages.
+
+ Attributes:
+ prompt_embeds: Optional serialized prompt embeddings payload for
+ direct transfer between stages
+ additional_information: Optional serialized additional information
+ dictionary containing tensors or lists to pass along with the request
+ """
+
# Optional prompt embeddings (direct-transfer version)
prompt_embeds: Optional[PromptEmbedsPayload] = None
# Optional additional information dictionary (serialized)
additional_information: Optional[AdditionalInformationPayload] = None
+
+
+class OmniEngineCoreOutput(
+ msgspec.Struct,
+ array_like=True, # type: ignore[call-arg]
+ omit_defaults=True, # type: ignore[call-arg]
+ gc=False,
+): # type: ignore[call-arg]
+ request_id: str
+ new_token_ids: list[int]
+
+ new_logprobs: Optional[LogprobsLists] = None
+ new_prompt_logprobs_tensors: Optional[LogprobsTensors] = None
+
+ pooling_output: Optional[dict[str, torch.Tensor]] = None
+
+ finish_reason: Optional[FinishReason] = None
+ stop_reason: Union[int, str, None] = None
+ events: Optional[list[EngineCoreEvent]] = None
+ kv_transfer_params: Optional[dict[str, Any]] = None
+
+ trace_headers: Optional[Mapping[str, str]] = None
+ # The number of tokens with prefix cache hits.
+ num_cached_tokens: int = 0
+
+ @property
+ def finished(self) -> bool:
+ return self.finish_reason is not None
+
+
+class OmniEngineCoreOutputs(
+ msgspec.Struct,
+ array_like=True, # type: ignore[call-arg]
+ omit_defaults=True, # type: ignore[call-arg]
+ gc=False,
+): # type: ignore[call-arg]
+ # NOTE(Nick): We could consider ways to make this more compact,
+ # e.g. columnwise layout
+
+ engine_index: int = 0
+
+ # [num_reqs]
+ outputs: list[OmniEngineCoreOutput] = []
+ scheduler_stats: Optional[SchedulerStats] = None
+ timestamp: float = 0.0
+
+ utility_output: Optional[UtilityOutput] = None
+ finished_requests: Optional[set[str]] = None
+
+ # In DP case, used to signal that the current wave of requests
+ # has finished and the engines are paused.
+ wave_complete: Optional[int] = None
+ # In DP case, used to signal that a request was received for an
+ # "old" wave, so the next wave needs to be started in other engines.
+ start_wave: Optional[int] = None
+
+ def __post_init__(self):
+ if self.timestamp == 0.0:
+ self.timestamp = time.monotonic()
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 64df79c7b..fc6dd5d7b 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -1,20 +1,50 @@
from dataclasses import dataclass
from typing import Optional
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeTextConfig
from vllm.engine.arg_utils import EngineArgs
+from vllm.logger import init_logger
from vllm.v1.engine.async_llm import AsyncEngineArgs
from vllm_omni.config import OmniModelConfig
+logger = init_logger(__name__)
+
@dataclass
class OmniEngineArgs(EngineArgs):
+ """Engine arguments for omni models, extending base EngineArgs.
+ Adds omni-specific configuration fields for multi-stage pipeline
+ processing and output type specification.
+ Args:
+ stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
+ model_stage: Stage type identifier, e.g., "thinker" or "talker"
+ (default: "thinker")
+ model_arch: Model architecture name
+ (default: "Qwen2_5OmniForConditionalGeneration")
+ engine_output_type: Optional output type specification for the engine.
+ Used to route outputs to appropriate processors (e.g., "image",
+ "audio", "latents"). If None, output type is inferred.
+ """
+
stage_id: int = 0
model_stage: str = "thinker"
model_arch: str = "Qwen2_5OmniForConditionalGeneration"
engine_output_type: Optional[str] = None
+ hf_config_name: Optional[str] = None
+
+ def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig:
+ # transformers' get_text_config method is used to get the text config from thinker_config.
+ # to handle the case that each model stage has their own text config,
+ # we need to draw the text config from the corresponding model stage.
+ return getattr(config_dict["hf_config"], config_dict["hf_config_name"]).get_text_config()
def create_model_config(self) -> OmniModelConfig:
+ """Create an OmniModelConfig from these engine arguments.
+ Returns:
+ OmniModelConfig instance with all configuration fields set
+ """
+
# First, get the base ModelConfig from the parent class
base_config = super().create_model_config()
@@ -27,7 +57,9 @@ class OmniEngineArgs(EngineArgs):
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["hf_config_name"] = self.hf_config_name
+ if self.hf_config_name is not None:
+ config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict)
# Create and return the OmniModelConfig instance
omni_config = OmniModelConfig(**config_dict)
omni_config.hf_config.architectures = omni_config.architectures
@@ -37,10 +69,31 @@ class OmniEngineArgs(EngineArgs):
@dataclass
class AsyncOmniEngineArgs(AsyncEngineArgs):
+ """Async engine arguments for omni models, extending base AsyncEngineArgs.
+ Adds omni-specific configuration fields for multi-stage pipeline
+ processing and output type specification in async contexts.
+ Args:
+ stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
+ model_stage: Stage type identifier, e.g., "thinker" or "talker"
+ (default: "thinker")
+ model_arch: Model architecture name
+ (default: "Qwen2_5OmniForConditionalGeneration")
+ engine_output_type: Optional output type specification for the engine.
+ Used to route outputs to appropriate processors (e.g., "image",
+ "audio", "latents"). If None, output type is inferred.
+ """
+
stage_id: int = 0
model_stage: str = "thinker"
model_arch: str = "Qwen2_5OmniForConditionalGeneration"
engine_output_type: Optional[str] = None
+ hf_config_name: Optional[str] = None
+
+ def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig:
+ # transformers' get_text_config method is used to get the text config from thinker_config.
+ # to handle the case that each model stage has their own text config,
+ # we need to draw the text config from the corresponding model stage.
+ return getattr(config_dict["hf_config"], config_dict["hf_config_name"]).get_text_config()
def create_model_config(self) -> OmniModelConfig:
# First, get the base ModelConfig from the parent class
@@ -56,6 +109,9 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
config_dict["model_arch"] = self.model_arch
config_dict["engine_output_type"] = self.engine_output_type
+ config_dict["hf_config_name"] = self.hf_config_name
+ if self.hf_config_name is not None:
+ config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict)
# Create and return the OmniModelConfig instance
omni_config = OmniModelConfig(**config_dict)
omni_config.hf_config.architectures = omni_config.architectures
diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py
index fb664b7a2..63e5ecdab 100644
--- a/vllm_omni/engine/output_processor.py
+++ b/vllm_omni/engine/output_processor.py
@@ -1,3 +1,4 @@
+from ast import Dict
from typing import Any, Callable, Optional, Union
import torch
@@ -19,6 +20,13 @@ logger = init_logger(__name__)
class OmniRequestState(RequestState):
+ """Request state for omni models, tracking multimodal outputs.
+
+ Extends the base RequestState with support for accumulating
+ multimodal tensor outputs (e.g., images, audio, latents) that
+ are produced incrementally during generation.
+ """
+
def __init__(
self,
*args,
@@ -26,7 +34,7 @@ class OmniRequestState(RequestState):
):
super().__init__(*args, **kwargs)
self.mm_type: Optional[str] = None
- self.mm_accumulated: Optional[torch.Tensor] = None
+ self.mm_accumulated: Optional[Dict[str, Any]] = None
@classmethod
def from_new_request(
@@ -85,22 +93,71 @@ class OmniRequestState(RequestState):
log_stats=log_stats,
)
- def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], mm_type: Optional[str]) -> None:
- if tensor is None:
+ def add_multimodal_tensor(self, payload: Optional[Any], mm_type: Optional[str]) -> None:
+ if payload 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
+
+ # Normalize incoming payload to dict on CPU
+ def _to_cpu(x):
+ if isinstance(x, torch.Tensor):
+ try:
+ return x.detach().to("cpu").contiguous()
+ except Exception:
+ return x
+ return x
+
+ if isinstance(payload, dict):
+ incoming: Dict[str, Any] = {}
+ # Optional remap: if producer used "model_outputs" or "hidden", rename to mm_type
+ # to keep a consistent key namespace per engine_core_output_type.
+ remapped_tensor = dict(payload)
+ target_key = self.mm_type or "hidden"
+ if "model_outputs" in remapped_tensor:
+ remapped_tensor[target_key] = remapped_tensor.pop("model_outputs")
+ elif "hidden" in remapped_tensor and target_key != "hidden":
+ remapped_tensor[target_key] = remapped_tensor.pop("hidden")
+ for k, v in remapped_tensor.items():
+ if isinstance(v, dict):
+ incoming[k] = {str(sk): _to_cpu(sv) for sk, sv in v.items()}
+ else:
+ incoming[k] = _to_cpu(v)
else:
- self.mm_accumulated = torch.cat([self.mm_accumulated, t], dim=0)
+ key = self.mm_type or "hidden"
+ incoming = {key: _to_cpu(payload)}
+
+ if self.mm_accumulated is None:
+ self.mm_accumulated = incoming
+ else:
+ # Merge keys; concatenate tensors along token dim when possible
+ for k, v in incoming.items():
+ if k not in self.mm_accumulated:
+ self.mm_accumulated[k] = v
+ else:
+ existing = self.mm_accumulated[k]
+ if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
+ try:
+ self.mm_accumulated[k] = torch.cat([existing, v], dim=0) # type: ignore[index]
+ except Exception:
+ self.mm_accumulated[k] = v
+ elif isinstance(v, dict) and isinstance(existing, dict):
+ # Merge nested dicts by concatenating tensors along token dim when possible
+ for sk, sv in v.items():
+ if sk not in existing:
+ existing[sk] = sv
+ continue
+ ev = existing[sk]
+ if isinstance(sv, torch.Tensor) and isinstance(ev, torch.Tensor):
+ try:
+ existing[sk] = torch.cat([ev, sv], dim=0)
+ except Exception:
+ existing[sk] = sv
+ else:
+ existing[sk] = sv
+ else:
+ self.mm_accumulated[k] = v
except Exception:
# Log and continue without crashing the output pipeline
logger.exception("Error accumulating multimodal tensor")
@@ -115,6 +172,23 @@ class OmniRequestState(RequestState):
stop_reason: Optional[Union[int, str]],
kv_transfer_params: Optional[dict[str, Any]] = None,
) -> Optional[Union[OmniRequestOutput, PoolingRequestOutput]]:
+ """Create a request output from generation results.
+
+ Creates a RequestOutput or PoolingRequestOutput from the generated
+ tokens and accumulated multimodal outputs. Attaches multimodal
+ tensors to the completion output if available.
+
+ Args:
+ new_token_ids: List of newly generated token IDs
+ pooling_output: Optional pooling output tensor
+ finish_reason: Optional finish reason indicating why generation stopped
+ stop_reason: Optional stop reason (token ID or stop string)
+ kv_transfer_params: Optional KV cache transfer parameters
+
+ Returns:
+ OmniRequestOutput or PoolingRequestOutput if output should be
+ emitted (based on finish status and output kind), None otherwise
+ """
finished = finish_reason is not None
final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
@@ -143,18 +217,15 @@ class OmniRequestState(RequestState):
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.
+ # Attach accumulated multimodal dict on the completion output
if not hasattr(base_output, "multimodal_output"):
setattr(base_output, "multimodal_output", {})
- setattr(base_output, "multimodal_output", {self.mm_type: tensor})
+ mm_out = getattr(base_output, "multimodal_output")
+ if isinstance(mm_out, dict):
+ for k, v in self.mm_accumulated.items():
+ mm_out[k] = v
+ else:
+ setattr(base_output, "multimodal_output", self.mm_accumulated)
except Exception:
logger.exception("Error in _new_completion_output")
return base_output
@@ -179,6 +250,15 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
log_stats: bool,
engine_core_output_type: Optional[str] = None,
):
+ """Initialize the multimodal output processor.
+
+ Args:
+ tokenizer: Tokenizer for detokenizing text outputs
+ log_stats: Whether to log statistics
+ engine_core_output_type: Optional output type specification
+ (e.g., "image", "audio", "latents"). Used to route outputs
+ to appropriate processors. If None, output type is inferred.
+ """
super().__init__(tokenizer=tokenizer, log_stats=log_stats)
self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
self._reqid_to_mm_type: dict[str, str] = {}
@@ -186,6 +266,16 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
self.engine_core_output_type = engine_core_output_type
def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None:
+ """Register a custom handler for a specific modality.
+
+ Allows custom processing logic for specific output modalities.
+ The handler is called before default processing for outputs
+ matching the specified modality.
+
+ Args:
+ modality: Modality name (e.g., "image", "audio", "latents")
+ handler: Callable that takes an EngineCoreOutput and processes it
+ """
self.output_handlers[modality.lower()] = handler
def add_request(
@@ -196,6 +286,21 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
request_index: int = 0,
queue: Optional[RequestOutputCollector] = None,
) -> None:
+ """Add a new request to be processed.
+
+ Creates an OmniRequestState for the request and registers it
+ for output processing.
+
+ Args:
+ request: Engine core request to add
+ prompt: Optional prompt string for the request
+ parent_req: Optional parent request for parallel sampling
+ request_index: Index of the request in the batch
+ queue: Optional queue for collecting outputs
+
+ Raises:
+ ValueError: If the request ID is already registered
+ """
request_id = request.request_id
if request_id in self.request_states:
raise ValueError(f"Request id {request_id} already running.")
@@ -220,6 +325,21 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
engine_core_timestamp: Optional[float] = None,
iteration_stats: Optional[IterationStats] = None,
) -> OutputProcessorOutput:
+ """Process engine core outputs into request outputs.
+
+ Converts EngineCoreOutput objects into RequestOutput objects,
+ handling multimodal outputs by routing them through appropriate
+ processors and accumulating tensors in request states.
+
+ Args:
+ engine_core_outputs: List of engine core outputs to process
+ engine_core_timestamp: Optional timestamp for the outputs
+ iteration_stats: Optional iteration statistics
+
+ Returns:
+ OutputProcessorOutput containing processed request outputs
+ and list of request IDs to abort
+ """
self._reqid_to_mm_type.clear()
for eco in engine_core_outputs:
mm_type = (self.engine_core_output_type or "").lower()
@@ -285,10 +405,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
# 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
+ ro.multimodal_output = req_state.mm_accumulated
except Exception:
logger.exception("Error attaching multimodal payload in process_outputs")
if req_state.queue is not None:
diff --git a/vllm_omni/engine/processor.py b/vllm_omni/engine/processor.py
index c25d7da0b..a93c98b7c 100644
--- a/vllm_omni/engine/processor.py
+++ b/vllm_omni/engine/processor.py
@@ -30,8 +30,28 @@ logger = init_logger(__name__)
class OmniProcessor(Processor):
+ """Processor for omni models, handling multimodal inputs and embeddings.
+
+ Extends the base vLLM Processor with support for processing prompt
+ embeddings and additional information payloads, enabling direct transfer
+ of pre-computed embeddings between pipeline stages.
+
+ Args:
+ vllm_config: Global vLLM configuration
+ tokenizer: Tokenizer instance for text processing
+ mm_registry: Multi-modal registry for processing multimodal inputs
+ """
+
@staticmethod
def _dtype_to_name(dtype: torch.dtype) -> str:
+ """Convert torch dtype to string representation.
+
+ Args:
+ dtype: PyTorch dtype to convert
+
+ Returns:
+ String representation of the dtype (e.g., "float32", "int64")
+ """
mapping = {
torch.float32: "float32",
torch.float: "float32",
@@ -78,6 +98,35 @@ class OmniProcessor(Processor):
priority: int = 0,
data_parallel_rank: Optional[int] = None,
) -> tuple[Optional[str], OmniEngineCoreRequest]:
+ """Process input prompt into an engine core request.
+
+ Converts a prompt (text, tokens, or multimodal) into an
+ OmniEngineCoreRequest that can be processed by the engine.
+ Handles prompt embeddings and additional information payloads
+ for direct transfer between stages.
+
+ Args:
+ request_id: Unique identifier for this request
+ prompt: Input prompt (text, token IDs, embeddings, or multimodal)
+ params: Sampling or pooling parameters for generation
+ arrival_time: Optional arrival timestamp (defaults to current time)
+ lora_request: Optional LoRA adapter request
+ tokenization_kwargs: Optional additional tokenization arguments
+ trace_headers: Optional tracing headers for observability
+ priority: Request priority (higher values processed first)
+ data_parallel_rank: Optional data parallel rank for distributed
+ inference
+
+ Returns:
+ Tuple of (prompt_string, OmniEngineCoreRequest) where:
+ - prompt_string: The original prompt as a string, or None if
+ using embeddings
+ - OmniEngineCoreRequest: Processed request ready for the engine
+
+ Raises:
+ ValueError: If data_parallel_rank is out of range or prompt_embeds
+ has incorrect shape
+ """
# TODO(woosuk): Support pooling models.
# TODO(woosuk): Support encoder-decoder models.
self._validate_lora(lora_request)
diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py
index cfde9e14a..5046196d0 100644
--- a/vllm_omni/entrypoints/async_omni_llm.py
+++ b/vllm_omni/entrypoints/async_omni_llm.py
@@ -2,7 +2,6 @@ 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
@@ -53,6 +52,38 @@ logger = init_logger(__name__)
class AsyncOmniLLM(EngineClient):
+ """Async entry point for vLLM-Omni inference.
+
+ This class provides an asynchronous interface for running multi-modal
+ comprehension and generation models. It orchestrates multiple
+ stages in a pipeline, where each stage runs in a separate process.
+ Designed for use with async/await patterns and streaming generation.
+
+ Args:
+ model: Model name or path to load
+ stage_configs_path: Optional path to YAML file containing stage
+ configurations. If None, configurations are loaded from the model.
+ log_stats: Whether to enable statistics logging
+ log_file: Optional path prefix for log files. If provided, logs will
+ be written to files with stage-specific suffixes.
+ init_sleep_seconds: Number of seconds to sleep between starting
+ each stage process during initialization
+ shm_threshold_bytes: Threshold in bytes for using shared memory
+ for IPC. Objects larger than this threshold will use shared memory.
+ batch_timeout: Timeout in seconds for batching requests within a stage
+ init_timeout: Timeout in seconds for waiting for all stages to initialize
+ **kwargs: Additional keyword arguments passed to stage engines
+
+ Example:
+ >>> async_llm = AsyncOmniLLM(model="Qwen/Qwen2.5-Omni-7B")
+ >>> async for output in async_llm.generate(
+ ... prompt="Hello",
+ ... request_id="req-1",
+ ... sampling_params_list=[SamplingParams(), SamplingParams()]
+ ... ):
+ ... print(output)
+ """
+
def __init__(
self,
model: str,
@@ -62,7 +93,7 @@ class AsyncOmniLLM(EngineClient):
init_sleep_seconds: int = 20,
shm_threshold_bytes: int = 65536,
batch_timeout: int = 10,
- init_timeout: int = 300,
+ init_timeout: int = 60000,
**kwargs,
):
self.batch_timeout = batch_timeout
@@ -78,7 +109,7 @@ class AsyncOmniLLM(EngineClient):
# Optional file handler for orchestrator
self._log_file = log_file
if self._log_file:
- remove_old_logs(self._log_file, len(self.stage_list))
+ remove_old_logs(self._log_file, len(self.stage_configs))
configure_orchestrator_logger(logger, self._log_file)
self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
@@ -140,6 +171,12 @@ class AsyncOmniLLM(EngineClient):
time.sleep(self._init_sleep_seconds)
def close(self) -> None:
+ """Close all stage processes and clean up resources.
+
+ Sends shutdown signals to all stage input queues and stops
+ all stage worker processes. This method should be called
+ when done using the AsyncOmniLLM instance.
+ """
for q in self._stage_in_queues:
try:
q.put_nowait(None)
@@ -156,13 +193,19 @@ class AsyncOmniLLM(EngineClient):
logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
def __del__(self) -> None: # best-effort
+ print("[AsyncOmniLLM] __del__ close()", flush=True)
+ raise Exception("test")
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."""
+ """Shutdown, cleaning up the background proc and IPC.
+
+ Alias for close() method. Cleans up all stage processes
+ and inter-process communication resources.
+ """
try:
self.close()
except Exception as e:
@@ -178,6 +221,33 @@ class AsyncOmniLLM(EngineClient):
priority: int = 0,
data_parallel_rank: Optional[int] = None,
) -> AsyncGenerator[OmniRequestOutput, None]:
+ """Generate outputs for the given prompt asynchronously.
+
+ Processes the prompt through all stages in the pipeline and yields
+ outputs as they become available. Each stage uses its corresponding
+ sampling parameters from the sampling_params_list.
+
+ Args:
+ prompt: Prompt to process. Can be a text string, token IDs,
+ or multimodal prompt.
+ request_id: Unique identifier for this request
+ sampling_params_list: List of SamplingParams, one for each stage.
+ Must have the same length as the number of stages.
+ If None, uses default sampling params for each stage.
+ lora_request: Optional LoRA adapter request for this generation
+ trace_headers: Optional tracing headers for observability
+ priority: Request priority (higher values processed first)
+ data_parallel_rank: Optional data parallel rank for distributed
+ inference
+
+ Yields:
+ OmniRequestOutput objects as they are produced by each stage.
+ Each output contains the stage_id, final_output_type, and
+ the request_output from that stage.
+
+ Raises:
+ ValueError: If sampling_params_list has incorrect length.
+ """
logger.debug("[Orchestrator] generate() called")
if sampling_params_list is None:
sampling_params_list = self.default_sampling_params_list
@@ -391,9 +461,8 @@ class AsyncOmniLLM(EngineClient):
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:
+ while len(self._stages_ready) < num_stages:
progressed = False
for stage_id, stage in enumerate(self.stage_list):
if stage_id in self._stages_ready:
@@ -459,20 +528,6 @@ class AsyncOmniLLM(EngineClient):
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.
@@ -510,9 +565,6 @@ class AsyncOmniLLM(EngineClient):
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:
@@ -569,6 +621,31 @@ class AsyncOmniLLM(EngineClient):
class AsyncOmniStageLLM(AsyncLLM):
+ """Async single-stage LLM engine for use within a stage worker process.
+
+ This class extends the base vLLM AsyncLLM class with omni-specific
+ processors for handling multimodal inputs and outputs. It is used
+ internally by AsyncOmniStage workers and should not be instantiated
+ directly by users.
+
+ Args:
+ engine_args: AsyncOmniEngineArgs containing engine configuration
+ vllm_config: Global vLLM configuration
+ executor_class: Executor implementation class, e.g. MultiprocExecutor
+ log_stats: Whether to log statistics
+ usage_context: Usage context of the LLM (default: ENGINE_CONTEXT)
+ mm_registry: Multi-modal registry for processing multimodal inputs
+ use_cached_outputs: Whether to use cached outputs
+ log_requests: Whether to log requests
+ start_engine_loop: Whether to start the engine loop automatically
+ stat_loggers: Customized stat loggers for the engine.
+ If not provided, default stat loggers will be used.
+ Note: Stat logger interface may change in V1.
+ client_addresses: Optional dictionary mapping client names to addresses
+ client_count: Total number of clients (default: 1)
+ client_index: Index of this client (default: 0)
+ """
+
def __init__(
self,
engine_args: AsyncOmniEngineArgs,
diff --git a/vllm_omni/entrypoints/cli/__init__.py b/vllm_omni/entrypoints/cli/__init__.py
index 2df4defa7..b233a71e6 100644
--- a/vllm_omni/entrypoints/cli/__init__.py
+++ b/vllm_omni/entrypoints/cli/__init__.py
@@ -1,4 +1,4 @@
-"""CLI helpers for vLLM-omni entrypoints."""
+"""CLI helpers for vLLM-Omni entrypoints."""
from .serve import OmniServeCommand
diff --git a/vllm_omni/entrypoints/cli/main.py b/vllm_omni/entrypoints/cli/main.py
index 148839d75..9647484f8 100644
--- a/vllm_omni/entrypoints/cli/main.py
+++ b/vllm_omni/entrypoints/cli/main.py
@@ -1,5 +1,5 @@
"""
-CLI entry point for vLLM-omni that intercepts vLLM commands.
+CLI entry point for vLLM-Omni that intercepts vLLM commands.
"""
import importlib.metadata
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index 25746cb5e..4606029ab 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -1,5 +1,5 @@
"""
-Omni serve command for vLLM-omni.
+Omni serve command for vLLM-Omni.
"""
import argparse
@@ -52,7 +52,7 @@ class OmniServeCommand(CLISubcommand):
serve_parser.add_argument(
"--omni",
action="store_true",
- help="Enable vLLM-omni mode",
+ help="Enable vLLM-Omni mode",
)
serve_parser.add_argument(
"--stage-configs-path",
diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py
new file mode 100644
index 000000000..5efc001da
--- /dev/null
+++ b/vllm_omni/entrypoints/omni.py
@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.diffusion.omni_diffusion import OmniDiffusion
+from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
+from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+
+class Omni:
+ """Unified entrypoint for both LLM and Diffusion models for better usability."""
+
+ def __init__(self, *args, **kwargs):
+ model = args[0] if args else kwargs.get("model", "")
+ if is_diffusion_model(model):
+ self.instance: OmniLLM | OmniDiffusion = OmniDiffusion(*args, **kwargs)
+ else:
+ self.instance: OmniLLM | OmniDiffusion = OmniLLM(*args, **kwargs)
+
+ def __getattr__(self, name):
+ """Delegate attribute access to the chosen backend instance."""
+ return getattr(self.instance, name)
+
+ def generate(self, *args, **kwargs):
+ """Convenience wrapper to call `generate` on the backend if available."""
+ if hasattr(self.instance, "generate"):
+ return getattr(self.instance, "generate")(*args, **kwargs)
+ raise AttributeError(f"'{self.instance.__class__.__name__}' has no attribute 'generate'")
diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py
index 655d0307f..4b9b1b4f6 100644
--- a/vllm_omni/entrypoints/omni_llm.py
+++ b/vllm_omni/entrypoints/omni_llm.py
@@ -1,6 +1,5 @@
import multiprocessing as mp
import os
-import sys
import time
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -11,7 +10,6 @@ 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
@@ -42,11 +40,36 @@ logger = init_logger(__name__)
class OmniLLM:
- """Multi-process pipelined OmniLLM.
+ """Main entry point for vLLM-Omni inference.
- - 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
+ This class provides a high-level interface for running multi-modal
+ comprehension and generation models. It orchestrates multiple
+ stages in a pipeline, where each stage runs in a separate process
+ with copy-based IPC (Queues). Downstream stages start when upstream
+ stages finish (window=-1), and each stage processes requests serially
+ (max_inflight=1) but pipelines across stages.
+
+ Args:
+ model: Model name or path to load
+ stage_configs_path: Optional path to YAML file containing stage
+ configurations. If None, configurations are loaded from the model.
+ log_stats: Whether to enable statistics logging
+ log_file: Optional path prefix for log files. If provided, logs will
+ be written to files with stage-specific suffixes.
+ init_sleep_seconds: Number of seconds to sleep between starting
+ each stage process during initialization
+ shm_threshold_bytes: Threshold in bytes for using shared memory
+ for IPC. Objects larger than this threshold will use shared memory.
+ batch_timeout: Timeout in seconds for batching requests within a stage
+ init_timeout: Timeout in seconds for waiting for all stages to initialize
+ **kwargs: Additional keyword arguments passed to stage engines
+
+ Example:
+ >>> llm = OmniLLM(model="Qwen/Qwen2.5-Omni-7B")
+ >>> outputs = llm.generate(
+ ... prompts="Hello",
+ ... sampling_params_list=[SamplingParams(), SamplingParams()]
+ ... )
"""
def __init__(
@@ -59,7 +82,7 @@ class OmniLLM:
shm_threshold_bytes: int = 65536,
batch_timeout: int = 10,
init_timeout: int = 300,
- **kwargs,
+ **kwargs: Any,
):
self.batch_timeout = batch_timeout
self._enable_stats: bool = bool(log_stats)
@@ -71,12 +94,12 @@ class OmniLLM:
# 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)
+ if self._log_file:
+ remove_old_logs(self._log_file, len(self.stage_list))
+ configure_orchestrator_logger(logger, self._log_file)
def _initialize_stages(
self,
@@ -132,6 +155,12 @@ class OmniLLM:
time.sleep(self._init_sleep_seconds)
def close(self) -> None:
+ """Close all stage processes and clean up resources.
+
+ Sends shutdown signals to all stage input queues and stops
+ all stage worker processes. This method should be called
+ when done using the OmniLLM instance.
+ """
for q in self._stage_in_queues:
try:
q.put_nowait(None)
@@ -157,6 +186,27 @@ class OmniLLM:
prompts: Union[PromptType, Sequence[PromptType]],
sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
) -> list[OmniRequestOutput]:
+ """Generate outputs for the given prompts.
+
+ Processes prompts through all stages in the pipeline and returns
+ the final outputs. Each stage uses its corresponding sampling
+ parameters from the sampling_params_list.
+
+ Args:
+ prompts: Single prompt or sequence of prompts to process.
+ Can be text strings, token IDs, or multimodal prompts.
+ sampling_params_list: List of SamplingParams, one for each stage.
+ Must have the same length as the number of stages.
+ Required for pipelined generation.
+
+ Returns:
+ List of OmniRequestOutput objects, one for each input prompt.
+ Each output contains the stage_id, final_output_type, and
+ the request_output from the final stage.
+
+ Raises:
+ ValueError: If sampling_params_list is None or has incorrect length.
+ """
try:
return self._run_generation(prompts, sampling_params_list)
except Exception as e:
@@ -326,7 +376,16 @@ class OmniLLM:
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]])
+ try:
+ next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
+ except Exception as e:
+ logger.exception(
+ "[Orchestrator] Process engine inputs error for req %s at stage %s: %s",
+ req_id,
+ next_stage_id,
+ e,
+ )
+ continue
sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index]
try:
# Measure transfer size and time (encode + enqueue)
@@ -415,7 +474,7 @@ class OmniLLM:
progressed = True
if result.get("type") == "stage_ready":
self._stages_ready.add(stage_id)
- logger.debug("[Orchestrator] Stage-%s reported ready", stage_id)
+ logger.info("[Orchestrator] Stage-%s reported ready", stage_id)
else:
# No user data should arrive before seeding; ignore other messages
pass
@@ -448,28 +507,33 @@ class OmniLLM:
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)
+ elif len(self._stages_ready) == num_stages:
+ logger.info("[Orchestrator] All stages initialized successfully")
class OmniStageLLM(LLM):
+ """Single-stage LLM engine for use within a stage worker process.
+
+ This class extends the base vLLM LLM class with omni-specific
+ processors for handling multimodal inputs and outputs. It is used
+ internally by OmniStage workers and should not be instantiated directly
+ by users.
+
+ Args:
+ model: Model name or path to load
+ compilation_config: Optional compilation configuration. Can be an
+ integer (compilation level), dict, or CompilationConfig instance.
+ hf_overrides: Optional HuggingFace model configuration overrides
+ structured_outputs_config: Optional structured outputs configuration.
+ Can be a dict or StructuredOutputsConfig instance.
+ **kwargs: Additional keyword arguments passed to the base LLM class
+ and engine
+ """
+
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,
):
@@ -500,9 +564,6 @@ class OmniStageLLM(LLM):
# 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)
@@ -527,7 +588,6 @@ class OmniStageLLM(LLM):
engine_args = OmniEngineArgs(
model=model,
- hf_overrides=hf_overrides,
compilation_config=compilation_config_instance,
structured_outputs_config=structured_outputs_instance,
**kwargs,
diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py
index 481c48213..d44c0043b 100644
--- a/vllm_omni/entrypoints/omni_stage.py
+++ b/vllm_omni/entrypoints/omni_stage.py
@@ -1,5 +1,5 @@
"""
-Stage manager for orchestrating multiple engines in vLLM-omni.
+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
@@ -14,6 +14,7 @@ import asyncio
import importlib
import logging
import multiprocessing as mp
+import os
from typing import Any, Optional, Union
from vllm.inputs import TextPrompt
@@ -39,7 +40,18 @@ logger = init_logger(__name__)
class OmniStage:
- def __init__(self, stage_config):
+ """Stage manager for orchestrating a single stage in the omni pipeline.
+
+ Encapsulates per-stage process lifecycle and worker logic, including
+ device setup, LLM initialization, batching, and shared-memory IPC.
+ Preserves input processing utilities for cross-stage data wiring.
+
+ Args:
+ stage_config: Stage configuration object containing engine arguments,
+ runtime settings, and stage-specific parameters
+ """
+
+ def __init__(self, stage_config: Any):
self.stage_config = stage_config
self.engine = None
self.async_engine = None
@@ -76,35 +88,69 @@ class OmniStage:
self._logger = logging.getLogger(__name__)
def set_engine(self, engine: LLMEngine) -> None:
- """Initialize the engine for the stage."""
+ """Set the LLM engine for this stage.
+
+ Args:
+ engine: LLMEngine instance to use for this stage
+ """
self.engine = engine
def set_async_engine(self, async_engine: AsyncLLM) -> None:
- """Initialize the async engine for the stage."""
+ """Set the async LLM engine for this stage.
+
+ Args:
+ async_engine: AsyncLLM instance to use for this 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)."""
+ def set_vllm_config(self, vllm_config: Any) -> None:
+ """Set the vLLM configuration for this stage.
+
+ Args:
+ vllm_config: VllmConfig instance 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)."""
+ """Set the tokenizer for this stage.
+
+ Args:
+ tokenizer: Tokenizer instance 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)."""
+ """Set the input preprocessor for this stage.
+
+ Args:
+ input_preprocessor: InputPreprocessor instance 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)."""
+ """Set whether tracing is enabled for this stage.
+
+ Args:
+ is_tracing_enabled: Boolean indicating if tracing is enabled
+ """
self.is_tracing_enabled = is_tracing_enabled
def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None:
- """Set the engine output for the stage."""
+ """Set the engine outputs for this stage.
+
+ Args:
+ engine_outputs: EngineCoreOutput from this stage's processing
+ """
self.engine_outputs = engine_outputs
# ----------------- New Orchestration APIs -----------------
def attach_queues(self, in_q: mp.Queue, out_q: mp.Queue) -> None:
+ """Attach input and output queues for IPC communication.
+
+ Args:
+ in_q: Input queue for receiving tasks from orchestrator
+ out_q: Output queue for sending results to orchestrator
+ """
self._in_q = in_q
self._out_q = out_q
@@ -118,6 +164,22 @@ class OmniStage:
ctx: Optional[mp.context.BaseContext] = None,
batch_timeout: int = 10,
) -> None:
+ """Initialize and start the stage worker process.
+
+ Creates a worker process that runs the LLM engine for this stage.
+ The worker handles batching, generation, and IPC communication.
+
+ Args:
+ model: Model name or path to load
+ is_async: Whether to use async engine (default: False)
+ log_file: Optional log file path prefix for stage-specific logs
+ shm_threshold_bytes: Threshold for using shared memory for IPC
+ ctx: Optional multiprocessing context (default: spawn)
+ batch_timeout: Timeout in seconds for batching requests
+
+ Raises:
+ AssertionError: If queues are not attached before calling this method
+ """
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
@@ -138,9 +200,6 @@ class OmniStage:
self,
model,
stage_payload,
- self._in_q,
- self._out_q,
- self._log_file,
batch_timeout,
),
)
@@ -159,6 +218,11 @@ class OmniStage:
self._proc.start()
def stop_stage_worker(self) -> None:
+ """Stop the stage worker process gracefully.
+
+ Sends shutdown signal to the worker and waits for it to terminate.
+ If graceful shutdown fails, forcefully terminates the process.
+ """
if self._in_q is not None:
try:
self._in_q.put_nowait(None)
@@ -174,12 +238,32 @@ class OmniStage:
self._proc.terminate()
except Exception as e:
self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
+ # Cleanup temporary stage log if we created one (only when no log_file provided)
+ try:
+ if not self._log_file:
+ tmp_path = f"/tmp/omni_stage{self.stage_id}.log"
+ if os.path.exists(tmp_path):
+ os.remove(tmp_path)
+ except Exception:
+ pass
def submit(self, payload: dict[str, Any]) -> None:
+ """Submit a task to the stage worker.
+
+ Args:
+ payload: Dictionary containing task data (request_id, engine_inputs,
+ sampling_params, etc.)
+ """
assert self._in_q is not None
self._in_q.put(payload)
def try_collect(self) -> Optional[dict[str, Any]]:
+ """Try to collect a result from the stage worker without blocking.
+
+ Returns:
+ Result dictionary if available, None otherwise. Result contains
+ request_id, engine_outputs (or engine_outputs_shm), and metrics.
+ """
assert self._out_q is not None
try:
return self._out_q.get_nowait()
@@ -187,17 +271,30 @@ class OmniStage:
return None
def process_engine_inputs(
- self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None
+ self, stage_list: list[Any], prompt: Union[OmniTokensPrompt, TextPrompt] = None
) -> list[Union[OmniTokensPrompt, TextPrompt]]:
- """Process the engine input for the stage."""
+ """Process engine inputs for this stage from upstream stage outputs.
+
+ Derives inputs for this stage from outputs of upstream stages.
+ Uses engine_input_source configuration to determine which upstream
+ stage outputs to use. Supports custom processing functions.
+
+ Args:
+ stage_list: List of all stages in the pipeline
+ prompt: Optional original prompt (for multimodal data preservation)
+
+ Returns:
+ List of processed engine inputs ready for this stage
+
+ Raises:
+ ValueError: If engine_input_source is empty or invalid
+ """
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)
@@ -232,6 +329,7 @@ def _stage_worker(
) -> None:
"""Stage worker entry: device setup, LLM init, batching, SHM IPC."""
import logging as _logging
+ import os as _os
import time as _time
from vllm_omni.entrypoints.log_utils import (
@@ -241,7 +339,6 @@ def _stage_worker(
log_stage_running_avg,
)
from vllm_omni.entrypoints.omni_llm import OmniStageLLM
-
# no inline JSONL/serialization imports; logging handled by utilities
stage_id = stage_payload["stage_id"]
@@ -249,12 +346,31 @@ def _stage_worker(
runtime_cfg = stage_payload.get("runtime", {})
shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
- # Per-stage file logger (optional)
+ # Per-stage logger: clear inherited handlers to avoid broken parent streams
try:
- if log_file:
- stage_log = _logging.getLogger(__name__)
- stage_log.setLevel(_logging.DEBUG)
- fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
+ stage_log = _logging.getLogger(__name__)
+ stage_log.setLevel(_logging.DEBUG)
+ for _h in list(stage_log.handlers):
+ try:
+ stage_log.removeHandler(_h)
+ except Exception:
+ pass
+ stage_log.propagate = False
+
+ class _StageFilter(_logging.Filter):
+ def filter(self, record: _logging.LogRecord) -> bool:
+ setattr(record, "stage", stage_id)
+ return True
+
+ # Prefer file logging even when log_file is not provided, to avoid invalid stdio in child procs
+ try:
+ if log_file:
+ _path = f"{log_file}.stage{stage_id}.log"
+ else:
+ _path = f"/tmp/omni_stage{stage_id}.log"
+ # Ensure dir exists
+ _os.makedirs(_os.path.dirname(_path), exist_ok=True)
+ fh = _logging.FileHandler(_path)
fh.setLevel(_logging.DEBUG)
fh.setFormatter(
_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")
@@ -267,6 +383,34 @@ def _stage_worker(
fh.addFilter(_StageFilter())
stage_log.addHandler(fh)
+ try:
+ # Also route vLLM internal logs to the same handler for this stage
+ _vllm_logger = _logging.getLogger("vllm")
+ _vllm_logger.setLevel(_logging.DEBUG)
+ for _vh in list(_vllm_logger.handlers):
+ try:
+ _vllm_logger.removeHandler(_vh)
+ except Exception:
+ pass
+ _vllm_logger.propagate = False
+ _vllm_logger.addHandler(fh)
+ except Exception:
+ pass
+ except Exception:
+ # Final fallback: attach NullHandler to avoid logging errors
+ stage_log.addHandler(_logging.NullHandler())
+ except Exception:
+ pass
+
+ try:
+ _logging.raiseExceptions = False
+ _root_logger = _logging.getLogger()
+ for _h in list(_root_logger.handlers):
+ try:
+ _root_logger.removeHandler(_h)
+ except Exception:
+ pass
+ _root_logger.addHandler(_logging.NullHandler())
except Exception:
pass
@@ -287,9 +431,7 @@ def _stage_worker(
# Init LLM
_logging.getLogger(__name__).debug(
- "[Stage-%s] Initializing engine with args keys=%s",
- stage_id,
- list(engine_args.keys()),
+ "[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)
@@ -304,21 +446,35 @@ def _stage_worker(
task = in_q.get()
_recv_dequeue_ts = _time.time()
if task is None:
- _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id)
+ _logging.getLogger(__name__).error("[Stage-%s] Received shutdown signal", stage_id)
break
max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
+ print(f"[Stage-{stage_id}] Max batch size: {max_batch_size}")
batch_tasks: list[dict[str, Any]] = [task]
if max_batch_size > 1:
while len(batch_tasks) < max_batch_size:
+ start_time = _time.time()
if not in_q.empty():
- extra = in_q.get(timeout=batch_timeout)
+ extra = in_q.get_nowait()
if extra is None:
in_q.put(None)
break
batch_tasks.append(extra)
+ end_time = _time.time()
+ duration = end_time - start_time
+ if duration > batch_timeout:
+ break
+ else:
+ continue
else:
- break
+ end_time = _time.time()
+ duration = end_time - start_time
+ _time.sleep(0.05)
+ if duration > batch_timeout:
+ break
+ else:
+ continue
batch_request_ids: list[Any] = []
batch_engine_inputs: list[Any] = []
@@ -344,7 +500,7 @@ def _stage_worker(
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))
+ _logging.getLogger(__name__).error("[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",
@@ -366,6 +522,14 @@ def _stage_worker(
gen_outputs.append(ro)
_gen_t1 = _time.time()
_gen_ms = (_gen_t1 - _gen_t0) * 1000.0
+ try:
+ print(
+ f"[Stage-{stage_id}] Generate done: batch={len(batch_tasks)}, "
+ f"req_ids={batch_request_ids}, gen_ms={_gen_ms:.1f}",
+ flush=True,
+ )
+ except Exception:
+ pass
# Group outputs per request id with fallback
req_to_outputs: dict[Any, list[Any]] = {rid: [] for rid in batch_request_ids}
@@ -487,21 +651,15 @@ 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))
+ asyncio.run(_stage_worker_async(omni_stage, model, stage_payload, 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."""
@@ -523,6 +681,9 @@ async def _stage_worker_async(
runtime_cfg = stage_payload.get("runtime", {})
shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
+ log_file = omni_stage._log_file
+ in_q = omni_stage._in_q
+ out_q = omni_stage._out_q
# Per-stage file logger (optional)
try:
if log_file:
@@ -696,6 +857,14 @@ async def _stage_worker_async(
"metrics": _metrics,
}
)
+ try:
+ print(
+ f"[Stage-{stage_id}] Enqueued req={rid}, use_shm={use_shm}, "
+ f"tokens_out={_metrics['num_tokens_out']}",
+ flush=True,
+ )
+ except Exception:
+ pass
except Exception as e:
_logging.getLogger(__name__).exception(
"[Stage-%s] Failed to enqueue result for request %s: %s",
@@ -728,3 +897,6 @@ async def _stage_worker_async(
"error": str(e),
}
)
+ print("--------------------------------", flush=True)
+ print(f"[Stage-{stage_id}] Stage worker exiting", flush=True)
+ print("--------------------------------", flush=True)
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index da403d804..61cec05b5 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -114,6 +114,21 @@ async def build_async_omni_llm(
disable_frontend_multiprocessing: Optional[bool] = None,
client_config: Optional[dict[str, Any]] = None,
) -> AsyncIterator[EngineClient]:
+ """Build an AsyncOmniLLM instance from command-line arguments.
+
+ Creates an async context manager that yields an AsyncOmniLLM instance
+ configured from the provided arguments. Handles forkserver setup if
+ needed and ensures proper cleanup on exit.
+
+ Args:
+ args: Parsed command-line arguments containing model and configuration
+ disable_frontend_multiprocessing: Optional flag to disable frontend
+ multiprocessing (deprecated in V1)
+ client_config: Optional client configuration dictionary
+
+ Yields:
+ EngineClient instance (AsyncOmniLLM) ready for use
+ """
if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver":
# The executor is expected to be mp.
# Pre-import heavy modules in the forkserver process
@@ -139,12 +154,23 @@ async def build_async_omni_llm_from_stage_config(
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
+ """Create AsyncOmniLLM from stage configuration.
- Returns the AsyncOmniLLM or None if the creation failed.
+ Creates an AsyncOmniLLM instance either in-process or using multiprocess
+ RPC. Loads stage configurations from the model or from a specified path.
+
+ Args:
+ args: Parsed command-line arguments containing model and stage configs
+ disable_frontend_multiprocessing: Flag to disable frontend multiprocessing
+ (deprecated in V1)
+ client_config: Optional client configuration dictionary
+
+ Yields:
+ EngineClient instance (AsyncOmniLLM) ready for use
+
+ Note:
+ Stage configurations are loaded from args.stage_configs_path if provided,
+ otherwise from the model's default configuration.
"""
# V1 AsyncLLM.
@@ -179,6 +205,17 @@ async def omni_init_app_state(
state: State,
args: Namespace,
) -> None:
+ """Initialize the FastAPI application state for omni API server.
+
+ Sets up the application state with model information, request logger,
+ and other server configuration needed for handling API requests.
+
+ Args:
+ engine_client: Engine client instance (AsyncOmniLLM)
+ vllm_config: vLLM configuration object
+ state: FastAPI application state object to initialize
+ args: Parsed command-line arguments
+ """
if args.served_model_name is not None:
served_model_names = args.served_model_name
else:
diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py
index 14c266c85..0cd522e14 100644
--- a/vllm_omni/entrypoints/openai/serving_chat.py
+++ b/vllm_omni/entrypoints/openai/serving_chat.py
@@ -730,7 +730,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
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()
+ audio_tensor = final_res.multimodal_output["audio"].float().detach().cpu().numpy()
# Convert numpy array to WAV bytes and encode as base64
if soundfile is None:
diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py
index 36ad2b0ab..8ddc7d504 100644
--- a/vllm_omni/entrypoints/stage_utils.py
+++ b/vllm_omni/entrypoints/stage_utils.py
@@ -53,11 +53,7 @@ def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
if 0 <= logical_idx < len(mapping):
selected_physical = mapping[logical_idx]
except Exception as e:
- logger.debug(
- "[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s",
- stage_id,
- e,
- )
+ logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e)
selected_physical = None
if selected_physical is None:
selected_physical = int(logical_idx)
@@ -69,22 +65,16 @@ def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
selected_physical,
)
elif devices in (None, "cpu"):
- logger.debug(
- "[Stage-%s] Using default device visibility (devices=%s)",
- stage_id,
- devices,
- )
+ logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
else:
selected_physical = int(str(devices))
os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
logger.debug(
- "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)",
- stage_id,
- selected_physical,
+ "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical
)
try:
- import torch # noqa: WPS433
+ import torch
if torch.cuda.is_available():
try:
diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py
index 7662304c0..857d9d491 100644
--- a/vllm_omni/entrypoints/utils.py
+++ b/vllm_omni/entrypoints/utils.py
@@ -13,8 +13,22 @@ logger = logging.getLogger(__name__)
PROJECT_ROOT = Path(__file__).parent.parent.parent
-def load_stage_configs_from_model(model: str):
- """Load stage configs from model."""
+def load_stage_configs_from_model(model: str) -> list:
+ """Load stage configurations from model's default config file.
+
+ Loads stage configurations based on the model type. Looks for a
+ YAML configuration file in the stage_configs directory matching
+ the model's model_type.
+
+ Args:
+ model: Model name or path (used to determine model_type)
+
+ Returns:
+ List of stage configuration dictionaries
+
+ Raises:
+ FileNotFoundError: If no stage config file exists for the model type
+ """
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"
@@ -25,7 +39,14 @@ def load_stage_configs_from_model(model: str):
return stage_configs
-def load_stage_configs_from_yaml(config_path: str):
- """Load stage configs from yaml file."""
+def load_stage_configs_from_yaml(config_path: str) -> list:
+ """Load stage configurations from a YAML file.
+
+ Args:
+ config_path: Path to the YAML configuration file
+
+ Returns:
+ List of stage configuration dictionaries from the file's stage_args
+ """
config_data = OmegaConf.load(config_path)
return config_data.stage_args
diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py
index 7c518bcc8..01a70d6a8 100644
--- a/vllm_omni/inputs/data.py
+++ b/vllm_omni/inputs/data.py
@@ -11,6 +11,17 @@ from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt
class OmniTokensPrompt(TokensPrompt):
+ """Tokens prompt with optional embeddings and additional information.
+
+ Extends TokensPrompt to support prompt embeddings and additional
+ information payloads for direct transfer between pipeline stages.
+
+ Attributes:
+ prompt_embeds: Optional tensor containing prompt embeddings
+ additional_information: Optional dictionary containing additional
+ information (tensors or lists) to pass along with the prompt
+ """
+
prompt_embeds: NotRequired[torch.Tensor]
"""The embeddings of the prompt."""
@@ -20,6 +31,18 @@ class OmniTokensPrompt(TokensPrompt):
class OmniTokenInputs(TokenInputs):
+ """Token inputs with optional embeddings and additional information.
+
+ Extends TokenInputs to support prompt embeddings and additional
+ information payloads for direct transfer between pipeline stages.
+
+ Attributes:
+ prompt_embeds: Optional tensor containing prompt embeddings
+ aligned with token IDs
+ additional_information: Optional dictionary containing additional
+ information (tensors or lists) to pass along with the inputs
+ """
+
# New: optional prompt embeddings aligned with token ids
prompt_embeds: NotRequired[torch.Tensor]
@@ -29,6 +52,17 @@ class OmniTokenInputs(TokenInputs):
class OmniEmbedsPrompt(EmbedsPrompt):
+ """Embeddings prompt with optional additional information.
+
+ Extends EmbedsPrompt to support additional information payloads
+ for direct transfer between pipeline stages.
+
+ Attributes:
+ prompt_embeds: Optional tensor containing prompt embeddings
+ additional_information: Optional dictionary containing additional
+ information (tensors or lists) to pass along with the prompt
+ """
+
# New: optional prompt embeddings aligned with token ids
prompt_embeds: NotRequired[torch.Tensor]
@@ -44,7 +78,22 @@ def token_inputs_omni(
prompt_embeds: Optional[torch.Tensor] = None,
additional_information: Optional[dict[str, Any]] = None,
) -> OmniTokenInputs:
- """Construct token inputs with optional embeddings and metadata."""
+ """Construct token inputs with optional embeddings and metadata.
+
+ Creates an OmniTokenInputs object with token IDs and optional
+ embeddings and additional information for pipeline stage transfer.
+
+ Args:
+ prompt_token_ids: List of token IDs for the prompt
+ prompt: Optional prompt string
+ cache_salt: Optional cache salt for prefix caching
+ prompt_embeds: Optional tensor containing prompt embeddings
+ additional_information: Optional dictionary containing additional
+ information (tensors or lists)
+
+ Returns:
+ OmniTokenInputs instance with the provided data
+ """
inputs = OmniTokenInputs(type="token", prompt_token_ids=prompt_token_ids)
if prompt is not None:
diff --git a/vllm_omni/inputs/parse.py b/vllm_omni/inputs/parse.py
index ace3b7a44..a3700e61f 100644
--- a/vllm_omni/inputs/parse.py
+++ b/vllm_omni/inputs/parse.py
@@ -9,6 +9,23 @@ from vllm.inputs.parse import (
def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPrompt:
+ """Parse a singleton prompt into a typed parsed prompt.
+
+ Handles omni-specific prompt types including tokens prompts with
+ embeddings and additional information. Supports string, text,
+ tokens, and embeddings prompts.
+
+ Args:
+ prompt: Singleton prompt to parse. Can be a string, TextPrompt,
+ TokensPrompt (with optional prompt_embeds and additional_information),
+ or EmbedsPrompt.
+
+ Returns:
+ ParsedSingletonPrompt containing the parsed prompt with type information
+
+ Raises:
+ TypeError: If the prompt type is not supported
+ """
if isinstance(prompt, str):
return ParsedStrPrompt(type="str", content=prompt)
elif isinstance(prompt, dict):
diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py
index 28046b802..3078d3457 100644
--- a/vllm_omni/inputs/preprocess.py
+++ b/vllm_omni/inputs/preprocess.py
@@ -13,6 +13,13 @@ logger = init_logger(__name__)
class OmniInputPreprocessor(InputPreprocessor):
+ """Input preprocessor for omni models.
+
+ Extends the base InputPreprocessor to handle omni-specific input
+ types including prompt embeddings and additional information payloads.
+ Supports processing tokens, embeddings, text, and multimodal inputs.
+ """
+
def _process_tokens(
self,
parsed_content: TokensPrompt,
diff --git a/vllm_omni/model_executor/layers/mrope.py b/vllm_omni/model_executor/layers/mrope.py
index 789b3b3cf..463653470 100644
--- a/vllm_omni/model_executor/layers/mrope.py
+++ b/vllm_omni/model_executor/layers/mrope.py
@@ -4,8 +4,11 @@ from typing import Optional, Union
import numpy as np
import torch
from transformers import PretrainedConfig
+from vllm.logger import init_logger
from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding
+logger = init_logger(__name__)
+
def _apply_rotary_emb(
x: torch.Tensor,
@@ -442,9 +445,16 @@ class MRotaryEmbedding(RotaryEmbedding):
# _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
+ try:
+ audio_token_id = thinker_config.audio_token_index
+ image_token_id = thinker_config.image_token_index
+ video_token_id = thinker_config.video_token_index
+ except Exception:
+ logger.info("Multimodal token idx changed!")
+ audio_token_id = thinker_config.audio_token_id
+ image_token_id = thinker_config.image_token_id
+ video_token_id = thinker_config.video_token_id
+
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
diff --git a/vllm_omni/model_executor/models/__init__.py b/vllm_omni/model_executor/models/__init__.py
index 46b0e61e0..0b2629b4a 100644
--- a/vllm_omni/model_executor/models/__init__.py
+++ b/vllm_omni/model_executor/models/__init__.py
@@ -1 +1,4 @@
+from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration
from .registry import OmniModelRegistry # noqa: F401
+
+__all__ = ["Qwen3OmniMoeForConditionalGeneration"]
diff --git a/vllm_omni/model_executor/models/output_templates.py b/vllm_omni/model_executor/models/output_templates.py
new file mode 100644
index 000000000..6fcd53eb3
--- /dev/null
+++ b/vllm_omni/model_executor/models/output_templates.py
@@ -0,0 +1,13 @@
+from typing import NamedTuple, Optional
+
+import torch
+from vllm.sequence import IntermediateTensors
+
+
+class OmniOutput(NamedTuple):
+ """Output from the merged Omni model containing both text and audio."""
+
+ text_hidden_states: torch.Tensor
+ multimodal_outputs: Optional[dict] = None
+ intermediate_tensors: Optional[IntermediateTensors] = None
+ next_token_id: Optional[torch.Tensor] = None
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
index b102b8ffb..e732777b4 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
@@ -2,7 +2,7 @@ import glob
import os
from collections.abc import Iterable
from functools import cached_property
-from typing import NamedTuple, Optional, Union
+from typing import Optional, Union
import numpy as np
import torch
@@ -26,6 +26,7 @@ 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.output_templates import OmniOutput
from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
Qwen2_5OmniConditionalGenerationMixin,
Qwen2_5OmniThinkerDummyInputsBuilder,
@@ -40,14 +41,6 @@ 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__)
@@ -304,7 +297,7 @@ class Qwen2_5OmniForConditionalGeneration(
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=None, multimodal_outputs={"model_outputs": audio_tensor})
return OmniOutput(
text_hidden_states=torch.cat(
diff --git a/vllm_omni/model_executor/models/qwen3_omni/__init__.py b/vllm_omni/model_executor/models/qwen3_omni/__init__.py
new file mode 100644
index 000000000..a469704f6
--- /dev/null
+++ b/vllm_omni/model_executor/models/qwen3_omni/__init__.py
@@ -0,0 +1,3 @@
+from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+
+__all__ = ["Qwen3OmniMoeForConditionalGeneration"]
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py
new file mode 100644
index 000000000..933236313
--- /dev/null
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py
@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.models.qwen3_moe import (
+ Qwen3MoeDecoderLayer,
+ Qwen3MoeMLP,
+ Qwen3MoeModel, # as _BaseQwen3MoeModel,
+)
+from vllm.model_executor.models.qwen3_moe import (
+ Qwen3MoeForCausalLM as _BaseQwen3MoeForCausalLM,
+)
+from vllm.model_executor.models.utils import (
+ PPMissingLayer,
+ maybe_prefix,
+)
+
+logger = init_logger(__name__)
+
+
+# Individual expert MoE block using Qwen3MoeMLP instead of FusedMoE
+class Qwen3OmniMoeSparseMoeBlock(nn.Module):
+ """Sparse MoE block using individual Qwen3MoeMLP experts instead of FusedMoE."""
+
+ def __init__(
+ self,
+ vllm_config: VllmConfig,
+ prefix: str = "",
+ ):
+ super().__init__()
+
+ config = vllm_config.model_config.hf_text_config
+ quant_config = vllm_config.quant_config
+
+ self.num_experts = config.num_experts
+ self.top_k = config.num_experts_per_tok
+ self.norm_topk_prob = config.norm_topk_prob
+ self.hidden_size = config.hidden_size
+
+ # Create individual expert MLPs
+ self.experts = nn.ModuleList(
+ [
+ Qwen3MoeMLP(
+ hidden_size=config.hidden_size,
+ intermediate_size=config.moe_intermediate_size,
+ hidden_act=config.hidden_act,
+ quant_config=quant_config,
+ reduce_results=False,
+ prefix=f"{prefix}.experts.{i}",
+ )
+ for i in range(self.num_experts)
+ ]
+ )
+
+ # Router for expert selection
+ from vllm.model_executor.layers.linear import ReplicatedLinear
+
+ self.gate = ReplicatedLinear(
+ config.hidden_size, config.num_experts, bias=False, quant_config=quant_config, prefix=f"{prefix}.gate"
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ """Forward pass using individual experts."""
+ # Handle 3D inputs (batch, seq_len, hidden_size) by reshaping to 2D
+ orig_shape = hidden_states.shape
+ if hidden_states.dim() == 3:
+ batch_size, seq_len, hidden_dim = hidden_states.shape
+ hidden_states = hidden_states.view(-1, hidden_dim)
+ elif hidden_states.dim() == 2:
+ num_tokens, hidden_dim = hidden_states.shape
+ elif hidden_states.dim() == 1:
+ hidden_states = hidden_states.unsqueeze(0)
+ num_tokens, hidden_dim = hidden_states.shape
+ else:
+ raise ValueError(
+ f"Qwen3OmniMoeSparseMoeBlock only supports 1D, 2D, or 3D inputs, got {hidden_states.dim()}D"
+ )
+
+ is_input_1d = len(orig_shape) == 1
+ hidden_states = hidden_states.view(-1, hidden_dim)
+
+ # Get router logits and select experts (matching transformers)
+ router_logits, _ = self.gate(hidden_states)
+ selected_experts, routing_weights = self._route_tokens(router_logits)
+
+ # Forward through individual experts
+ final_hidden_states = self._forward_experts(hidden_states, selected_experts, routing_weights)
+
+ # Reshape back to original shape
+ if is_input_1d:
+ return final_hidden_states.squeeze(0)
+ elif len(orig_shape) == 3:
+ # Reshape back to 3D (batch, seq_len, hidden_dim)
+ return final_hidden_states.view(orig_shape)
+ else:
+ return final_hidden_states
+
+ def _route_tokens(self, router_logits: torch.Tensor):
+ """Route tokens to experts using top-k selection (matching transformers)."""
+ routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
+ routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+ if self.norm_topk_prob:
+ routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+ routing_weights = routing_weights.to(router_logits.dtype)
+ return selected_experts, routing_weights
+
+ def _forward_experts(
+ self, hidden_states: torch.Tensor, selected_experts: torch.Tensor, routing_weights: torch.Tensor
+ ):
+ """Forward through individual experts (matching transformers implementation)."""
+ final_hidden_states = torch.zeros_like(hidden_states)
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+ expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+ for expert_idx in expert_hit:
+ idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
+ current_state = hidden_states[None, top_x].reshape(-1, hidden_states.shape[-1])
+ current_hidden_states = self.experts[expert_idx](current_state) * routing_weights[top_x, idx, None]
+ final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+
+ return final_hidden_states
+
+
+class Qwen3MoeForCausalLM(_BaseQwen3MoeForCausalLM):
+ """Thin wrapper to swap in the patched `Qwen3MoeModel`."""
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ # Don't call super().__init__() to avoid duplicate layer registration.
+ nn.Module.__init__(self)
+ config = vllm_config.model_config.hf_text_config
+ quant_config = vllm_config.quant_config
+ self.config = config
+ self.quant_config = quant_config
+ self.model = Qwen3MoeModel(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+ self.lm_head = ParallelLMHead(
+ config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=maybe_prefix(prefix, "lm_head")
+ )
+ if self.config.tie_word_embeddings:
+ self.lm_head.weight = self.model.embed_tokens.weight
+ self.logits_processor = LogitsProcessor(config.vocab_size)
+ self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
+
+ # Set MoE hyperparameters for individual experts
+ self.expert_weights = []
+
+ self.moe_layers: list[FusedMoE] = []
+ example_layer = None
+ for layer in self.model.layers:
+ if isinstance(layer, PPMissingLayer):
+ continue
+
+ assert isinstance(layer, Qwen3MoeDecoderLayer)
+ if isinstance(layer.mlp, FusedMoE):
+ example_layer = layer.mlp
+ self.moe_layers.append(layer.mlp)
+
+ if example_layer is None:
+ raise RuntimeError("No Qwen3OmniMoe layer found in the model.layers.")
+
+ self.num_moe_layers = len(self.moe_layers)
+ self.num_expert_groups = 1
+ self.num_shared_experts = 0
+ self.num_logical_experts = example_layer.n_logical_experts
+ self.num_physical_experts = example_layer.n_physical_experts
+ self.num_local_physical_experts = example_layer.n_local_physical_experts
+ self.num_routed_experts = example_layer.n_routed_experts
+ self.num_redundant_experts = example_layer.n_redundant_experts
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
new file mode 100644
index 000000000..2709028ad
--- /dev/null
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
@@ -0,0 +1,1132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright 2025 The Qwen team.
+"""Inference-only Qwen3-Omni-Moe unified model (thinker + talker + code2wav)."""
+
+from collections import defaultdict
+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.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+ Qwen3OmniMoeCode2WavConfig,
+ Qwen3OmniMoeConfig,
+ Qwen3OmniMoeTalkerConfig,
+ Qwen3OmniMoeThinkerConfig,
+)
+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.utils import init_vllm_registered_model, maybe_prefix
+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.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
+
+from .qwen3_omni_moe_thinker import (
+ Qwen3OmniMoeConditionalGenerationMixin,
+ Qwen3OmniMoeThinkerDummyInputsBuilder,
+ Qwen3OmniMoeThinkerMultiModalProcessor,
+ Qwen3OmniMoeThinkerProcessingInfo,
+)
+
+# Special token IDs for Qwen3 Omni MoE
+# Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
+
+# Audio tokens (thinker vocabulary, for marking audio boundaries)
+AUDIO_START_TOKEN_ID = 151669 # <|audio_start|> (audio_bos_token)
+AUDIO_END_TOKEN_ID = 151670 # <|audio_end|> (audio_eos_token)
+AUDIO_PAD_TOKEN_ID = 151675 # <|audio_pad|>
+
+# TTS text tokens (thinker vocabulary, for text-to-speech control)
+TTS_PAD_TOKEN_ID = 151671 # <tts_pad>
+TTS_BOS_TOKEN_ID = 151672 # <tts_text_bos>
+TTS_EOS_TOKEN_ID = 151673 # <tts_text_eod> (end of dialogue)
+TTS_BOS_SINGLE_TOKEN_ID = 151674 # <tts_text_bos_single>
+
+# Talker codec tokens (talker vocabulary, used for RVQ code generation)
+TALKER_CODEC_PAD_TOKEN_ID = 4196 # Padding token
+TALKER_CODEC_BOS_TOKEN_ID = 4197 # Beginning of speech
+TALKER_CODEC_EOS_TOKEN_ID = 4198 # End of speech
+TALKER_CODEC_NOTHINK_ID = 4203 # No-think mode
+TALKER_CODEC_THINK_BOS_ID = 4204 # Think mode start
+TALKER_CODEC_THINK_EOS_ID = 4205 # Think mode end
+
+logger = init_logger(__name__)
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+ Qwen3OmniMoeThinkerMultiModalProcessor,
+ info=Qwen3OmniMoeThinkerProcessingInfo,
+ dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder,
+)
+class Qwen3OmniMoeForConditionalGeneration(
+ nn.Module, SupportsMultiModal, SupportsPP, Qwen3OmniMoeConditionalGenerationMixin
+):
+ """
+ Unified Qwen3 Omni MoE model combining thinker, talker, and code2wav.
+
+ Architecture:
+ - Thinker: Multimodal understanding (text + audio + video) → text generation
+ - Talker: Text embeddings → RVQ codec codes
+ - Code2Wav: RVQ codes → audio waveform
+
+ Usage:
+ Set `model_stage` in vllm_config to one of: "thinker", "talker", "code2wav"
+ """
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ super().__init__()
+ self.have_multimodal_outputs = True
+ config: Qwen3OmniMoeConfig = 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
+ self.config = config
+
+ # Initialize thinker components
+ thinker_config: Qwen3OmniMoeThinkerConfig = config.thinker_config
+ self.thinker_config = thinker_config
+ self.multimodal_config = multimodal_config
+
+ # Initialize talker components
+ talker_config: Qwen3OmniMoeTalkerConfig = config.talker_config
+ self.talker_config = talker_config
+
+ # Initialize code2wav components
+ code2wav_config: Qwen3OmniMoeCode2WavConfig = config.code2wav_config
+ self.code2wav_config = code2wav_config
+
+ # Determine model stage
+ self.model_stage = vllm_config.model_config.model_stage
+
+ if self.model_stage == "thinker":
+ # Initialize thinker model (multimodal processing + text generation)
+ # Create a new vllm_config with thinker_config as the hf_config
+ thinker_vllm_config = vllm_config.with_hf_config(
+ thinker_config, architectures=["Qwen3OmniMoeThinkerForConditionalGeneration"]
+ )
+ self.thinker = init_vllm_registered_model(
+ vllm_config=thinker_vllm_config,
+ prefix=maybe_prefix(prefix, "thinker"),
+ hf_config=thinker_config,
+ architectures=["Qwen3OmniMoeThinkerForConditionalGeneration"],
+ )
+ self.model = self.thinker
+ self.talker = None
+ self.code2wav = None
+ elif self.model_stage == "talker":
+ self.thinker = None
+ # Initialize talker model (text embeddings → codec codes)
+ # Create a new vllm_config with talker_config as the hf_config
+ # This ensures the talker uses its own text_config (smaller vocab_size)
+ talker_vllm_config = vllm_config.with_hf_config(
+ talker_config, architectures=["Qwen3OmniMoeTalkerForConditionalGeneration"]
+ )
+ self.talker = init_vllm_registered_model(
+ vllm_config=talker_vllm_config,
+ prefix=maybe_prefix(prefix, "talker"),
+ hf_config=talker_config,
+ architectures=["Qwen3OmniMoeTalkerForConditionalGeneration"],
+ )
+ self.talker.init_multi_modal(thinker_config)
+ self.model = self.talker
+ self.code2wav = None
+ from transformers.generation.logits_process import (
+ LogitsProcessorList,
+ RepetitionPenaltyLogitsProcessor,
+ SuppressTokensLogitsProcessor,
+ TemperatureLogitsWarper,
+ TopKLogitsWarper,
+ )
+
+ self.logits_processors = LogitsProcessorList(
+ [
+ RepetitionPenaltyLogitsProcessor(penalty=1.05),
+ TemperatureLogitsWarper(temperature=0.9),
+ SuppressTokensLogitsProcessor(suppress_tokens=self._get_talker_suppressed_tokens()),
+ TopKLogitsWarper(top_k=50),
+ ]
+ )
+
+ elif self.model_stage == "code2wav":
+ self.thinker = None
+ self.talker = None
+ # Initialize code2wav (codec codes → audio waveform)
+ # Create a new vllm_config with code2wav_config as the hf_config
+ code2wav_vllm_config = vllm_config.with_hf_config(code2wav_config, architectures=["Qwen3OmniMoeCode2Wav"])
+ self.code2wav = init_vllm_registered_model(
+ vllm_config=code2wav_vllm_config,
+ prefix=maybe_prefix(prefix, "code2wav"),
+ hf_config=code2wav_config,
+ architectures=["Qwen3OmniMoeCode2Wav"],
+ )
+ self.model = self.code2wav
+ else:
+ raise ValueError(
+ f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'talker', 'code2wav'"
+ )
+
+ # 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:
+ """Get the device of a module."""
+ 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,
+ code_predictor_device: Optional[Union[str, torch.device]] = None,
+ code2wav_device: Optional[Union[str, torch.device]] = None,
+ ) -> None:
+ """
+ Optionally move thinker/talker/code2wav to different devices.
+
+ Example:
+ model.move_submodules_to_devices(
+ thinker_device='cuda:0',
+ talker_device='cuda:1',
+ code2wav_device='cuda:2',
+ )
+ """
+ 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 code2wav_device is not None and self.code2wav is not None:
+ self.code2wav.to(code2wav_device)
+
+ @cached_property
+ def sampler(self):
+ """Get sampler from active model."""
+ 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:
+ """Get input embeddings for the active model stage."""
+ if self.model_stage == "code2wav":
+ # Code2wav doesn't use text embeddings
+ 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 active model for multimodal processing."""
+ return self.model.get_multimodal_embeddings(**kwargs)
+
+ # ==================== Forward Pass ====================
+ def _get_talker_suppressed_tokens(self):
+ return [
+ i
+ for i in range(
+ self.config.talker_config.text_config.vocab_size - 1024,
+ self.config.talker_config.text_config.vocab_size,
+ )
+ if i != self.config.talker_config.codec_eos_token_id
+ ]
+
+ 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 = "ethan",
+ codec: Optional[torch.Tensor] = None,
+ sampling_metadata: Optional[SamplingMetadata] = None,
+ logits_index: Optional[int] = None,
+ additional_information: Optional[dict[str, object]] = None,
+ **kwargs: object,
+ ) -> Union[torch.Tensor, IntermediateTensors, OmniOutput]:
+ """
+ Unified forward pass for all model stages.
+
+ Workflow:
+ 1) Thinker: multimodal understanding → text hidden states
+ 2) Talker -> Code Predictor: text embeddings → codec codes (layer 0 + code_predictor:residual layers)
+ 3) Code2wav: 8-layer RVQ codes → audio waveform
+
+ Returns:
+ OmniOutput with text_hidden_states and optional audio
+ """
+
+ # ========== Stage 1: Thinker ==========
+ if self.model_stage == "thinker":
+ # Normalize to batched inputs if needed
+ 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)
+
+ # Handle None input_ids
+ if input_ids is None:
+ input_ids = torch.zeros(
+ inputs_embeds.shape[1],
+ dtype=torch.long,
+ device=thinker_dev,
+ ).unsqueeze(0)
+ added_batch_dim = True
+
+ # Move to thinker 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 forward
+ # If talker expects a specific intermediate layer, capture it here
+ accept_layer = getattr(self.talker_config, "accept_hidden_layer", None)
+ capture_kwargs = {}
+ if accept_layer is not None:
+ capture_kwargs = {
+ "capture_layer_indices": [0, int(accept_layer)],
+ "return_hidden_states": True,
+ }
+ text_hidden_states, captured_layer_dict = self.thinker(
+ input_ids=input_ids,
+ positions=positions[0] if positions.ndim > 1 else positions,
+ intermediate_tensors=intermediate_tensors,
+ inputs_embeds=inputs_embeds,
+ **capture_kwargs,
+ **kwargs,
+ )
+ # Compute thinker-side TTS token embeddings for BOS/EOS/PAD and expose via multimodal outputs.
+ # These will later be projected into talker text space by the talker stage.
+ multimodal_outputs = captured_layer_dict if captured_layer_dict is not None else {}
+ try:
+ tts_tokens = torch.tensor(
+ [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
+ device=self._module_device(self.thinker),
+ dtype=torch.long,
+ )
+ thinker_tts_embeds = self.thinker.get_input_embeddings(tts_tokens) # [1,3,thinker_hidden]
+ if (
+ isinstance(thinker_tts_embeds, torch.Tensor)
+ and thinker_tts_embeds.ndim == 3
+ and thinker_tts_embeds.shape[1] == 3
+ ):
+ bos_eos_pad = thinker_tts_embeds.to(text_hidden_states.device).chunk(3, dim=1) # 3 * [1,1,H]
+ multimodal_outputs["tts_bos_embed"] = bos_eos_pad[0]
+ multimodal_outputs["tts_eos_embed"] = bos_eos_pad[1]
+ multimodal_outputs["tts_pad_embed"] = bos_eos_pad[2]
+ except Exception:
+ # Best-effort; absence will be handled by talker with fallbacks
+ pass
+
+ # Return text-only output (with multimodal sidecar)
+ return OmniOutput(
+ text_hidden_states=(text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states),
+ multimodal_outputs=multimodal_outputs,
+ )
+
+ # ========== Stage 2.1: Talker ==========
+ elif 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.code_predictor_hidden_pool = 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., code_predictor_hidden_per_request)
+ update_by_req_id: dict[str, dict] = {}
+ text_step = torch.zeros(
+ 1,
+ self.talker_config.text_config.hidden_size,
+ device=self._module_device(self.talker),
+ dtype=torch.bfloat16,
+ )
+ last_talker_hidden = torch.zeros(
+ 1,
+ 1,
+ self.talker_config.text_config.hidden_size,
+ device=self._module_device(self.talker),
+ dtype=torch.bfloat16,
+ )
+
+ batched_talker_inputs = defaultdict(dict)
+ if is_profile:
+ batched_talker_inputs["0"] = {
+ "input_ids": input_ids,
+ "positions": positions,
+ "inputs_embeds": inputs_embeds.clone(),
+ "text_step": text_step,
+ "last_talker_hidden": last_talker_hidden,
+ }
+
+ # ------- 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
+ batched_talker_inputs[rid] = {
+ "input_ids": input_ids[s:e],
+ "positions": positions[s:e] if len(positions.shape) == 1 else positions[0, s:e],
+ "inputs_embeds": inputs_embeds[s:e].clone(),
+ "text_step": text_step,
+ "last_talker_hidden": last_talker_hidden,
+ }
+ info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {}
+ # Read thinker outputs for prefill
+ thinker_sequence_embeds = info.get("thinker_embeddings").to(
+ device=self._module_device(self.talker), dtype=torch.bfloat16
+ ) # Tensor [P,H]
+ thinker_hidden_states = info.get("thinker_hidden_states").to(
+ device=self._module_device(self.talker), dtype=torch.bfloat16
+ ) # Tensor [K,H]
+ thinker_sequences = (
+ info.get("thinker_sequences")
+ if info.get("thinker_sequences") is None
+ else torch.as_tensor(info.get("thinker_sequences"), device=self._module_device(self.talker))
+ )
+ thinker_chatml_ids = (
+ info.get("thinker_input_ids")
+ if info.get("thinker_input_ids") is None
+ else torch.as_tensor(info.get("thinker_input_ids"), device=self._module_device(self.talker))
+ )
+ # Optional thinker-side TTS token embeddings (to be projected by talker)
+ tts_bos_thinker = info.get("tts_bos_embed").to(
+ device=self._module_device(self.talker), dtype=torch.bfloat16
+ )
+ tts_eos_thinker = info.get("tts_eos_embed").to(
+ device=self._module_device(self.talker), dtype=torch.bfloat16
+ )
+ tts_pad_thinker = info.get("tts_pad_embed").to(
+ device=self._module_device(self.talker), dtype=torch.bfloat16
+ )
+
+ if thinker_sequence_embeds is None or thinker_hidden_states is None:
+ raise ValueError(
+ f"additional_information_by_req_id[{rid}] must include "
+ f"'thinker_embeddings' and 'thinker_hidden_states' for talker prefill."
+ )
+
+ # Normalize to tensors
+ if not isinstance(thinker_sequence_embeds, torch.Tensor):
+ thinker_sequence_embeds = torch.as_tensor(
+ thinker_sequence_embeds, device=self._module_device(self.talker)
+ )
+ if not isinstance(thinker_hidden_states, torch.Tensor):
+ thinker_hidden_states = torch.as_tensor(
+ thinker_hidden_states, device=self._module_device(self.talker)
+ )
+
+ if isinstance(thinker_chatml_ids, torch.Tensor) or isinstance(thinker_chatml_ids, list):
+ ids_chatml = (
+ thinker_chatml_ids
+ if isinstance(thinker_chatml_ids, torch.Tensor)
+ else torch.as_tensor(thinker_chatml_ids, device=self._module_device(self.talker))
+ )
+ if ids_chatml.ndim == 1:
+ ids_chatml = ids_chatml.unsqueeze(0)
+ else:
+ # Fallback: create dummy ids if not provided
+ ids_chatml = torch.zeros(
+ (1, thinker_sequence_embeds.shape[1]),
+ dtype=torch.long,
+ device=self._module_device(self.talker),
+ )
+ thinker_sequences = ids_chatml
+
+ speaker_id = self._get_text_spk_token_id(voice_type)
+ req_input_ids, req_embeds, trailing_text_hidden = self._thinker_to_talker_prefill(
+ thinker_embed=thinker_sequence_embeds.to(self._module_device(self.talker)),
+ thinker_hidden=thinker_hidden_states.to(self._module_device(self.talker)),
+ multimodal_mask=None,
+ input_ids=ids_chatml.to(self._module_device(self.talker)),
+ thinker_result_ids=thinker_sequences.to(self._module_device(self.talker)),
+ speaker_id=speaker_id,
+ tts_bos_thinker=tts_bos_thinker,
+ tts_eos_thinker=tts_eos_thinker,
+ tts_pad_thinker=tts_pad_thinker,
+ )
+ # req_embeds: [S,D], req_input_ids: [S] → map到当前 span
+ seg_len = min(span_len, req_embeds.shape[0])
+ # 强制对齐 dtype/device/内存布局，避免精度或布局差异导致比较不相等
+ inputs_embeds[s : s + seg_len] = (
+ req_embeds[:seg_len].to(device=inputs_embeds.device, dtype=inputs_embeds.dtype).contiguous()
+ )
+
+ if isinstance(req_input_ids, torch.Tensor) and req_input_ids.shape[0] >= seg_len:
+ input_ids[s : s + seg_len] = req_input_ids[:seg_len].to(input_ids.device)
+ # Queue tailing_text_hidden for decode (drop first for next steps),
+ # similar to Qwen2.5's thinker_reply_part_per_request
+ batched_talker_inputs[rid]["inputs_embeds"] = inputs_embeds[s : s + seg_len].clone()
+ try:
+ if isinstance(trailing_text_hidden, torch.Tensor) and trailing_text_hidden.numel() > 0:
+ if trailing_text_hidden.ndim == 2:
+ rem_tail = trailing_text_hidden
+ elif trailing_text_hidden.ndim == 1:
+ rem_tail = torch.zeros(
+ 0,
+ trailing_text_hidden.shape[0],
+ dtype=trailing_text_hidden.dtype,
+ device=trailing_text_hidden.device,
+ )
+ else:
+ # compatible with old shape [1,S,D]
+ rem_tail = trailing_text_hidden.squeeze(0)
+ if rem_tail.shape[0] > 0:
+ update_by_req_id.setdefault(rid, {})["tailing_text_hidden_per_request"] = (
+ rem_tail.detach().to("cpu").contiguous()
+ )
+ # Also persist projected tts_pad for decode fallback if needed
+ if isinstance(tts_pad_thinker, torch.Tensor):
+ pad_in = tts_pad_thinker
+ if pad_in.ndim == 2:
+ pad_in = pad_in.unsqueeze(0)
+ if pad_in.ndim == 1:
+ pad_in = pad_in.view(1, 1, -1)
+ pad_proj = self.talker.text_projection(pad_in.to(self._module_device(self.talker)))
+ update_by_req_id.setdefault(rid, {})["tts_pad_embed_projected"] = (
+ pad_proj.detach().to("cpu").contiguous()
+ )
+ except Exception:
+ pass
+
+ # ------- 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
+ batched_talker_inputs[rid] = {
+ "input_ids": input_ids[s:e],
+ "positions": positions[s:e] if len(positions.shape) == 1 else positions[0, s:e],
+ "inputs_embeds": inputs_embeds[s:e].clone(),
+ "text_step": text_step,
+ "last_talker_hidden": last_talker_hidden,
+ }
+ try:
+ q_tail = None
+ if isinstance(rid, str):
+ q_tail = runtime_addi_by_req.get(rid, {}).get("tailing_text_hidden_per_request")
+ if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0:
+ use_vec = q_tail[0:1, :]
+ new_q_tail = (
+ q_tail[1:, :].detach().to("cpu").contiguous()
+ if q_tail.shape[1] > 1
+ else self.tts_pad_embed.to(inputs_embeds.device, dtype=inputs_embeds.dtype)
+ )
+ text_step = use_vec.to(inputs_embeds.device, dtype=inputs_embeds.dtype)
+ update_by_req_id.setdefault(rid, {})["tailing_text_hidden_per_request"] = new_q_tail
+ else:
+ text_step = self.tts_pad_embed.to(inputs_embeds.device, dtype=inputs_embeds.dtype)
+
+ if isinstance(rid, str):
+ last_talker_hidden = (
+ runtime_addi_by_req.get(rid, {})
+ .get("last_talker_hidden")
+ .to(inputs_embeds.device, dtype=inputs_embeds.dtype)
+ )
+ last_talker_hidden = last_talker_hidden.reshape(
+ *last_talker_hidden.shape[-2:]
+ ) # [1, hidden_size]
+ batched_talker_inputs[rid]["last_talker_hidden"] = last_talker_hidden
+ batched_talker_inputs[rid]["text_step"] = text_step
+ except Exception as e:
+ logger.error(f"Error in decode: {e}")
+
+ if not is_profile and input_ids.shape[0] > 1:
+ self.input_ids_list = input_ids.tolist()
+
+ last_talker_hidden = None
+ # Run talker forward
+ with torch.inference_mode():
+ talker_hidden, batched_code_predictor_codes = self.talker(
+ batched_talker_inputs=batched_talker_inputs,
+ )
+ if not is_profile:
+ for rid in request_ids:
+ idx_req = request_ids.index(rid)
+ s, e = request_token_spans[idx_req]
+ span_len = int(e) - int(s)
+ update_by_req_id.setdefault(rid, {})["last_talker_hidden"] = (
+ talker_hidden[e - 1 : e, :].detach().to("cpu").contiguous()
+ )
+
+ 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 not is_profile and multimodal_outputs is not None and batched_code_predictor_codes is not None:
+ multimodal_outputs.update({"code_predictor_codes": batched_code_predictor_codes})
+ elif not is_profile and multimodal_outputs is None and batched_code_predictor_codes is not None:
+ multimodal_outputs = {"code_predictor_codes": batched_code_predictor_codes}
+
+ next_token_id = None
+
+ return OmniOutput(
+ text_hidden_states=talker_hidden,
+ multimodal_outputs=multimodal_outputs,
+ next_token_id=next_token_id,
+ )
+
+ # ========== Stage 3: Code2Wav ==========
+ elif self.model_stage == "code2wav":
+ # Extract codec codes from input
+ codes = []
+ if input_ids is not None:
+ # addi_by_req = kwargs.get("additional_information_by_req_id")
+ codes.append(input_ids.reshape(1, 16, -1))
+
+ else:
+ # for profile, we use max length from inputs_embeds
+ codes.append(
+ torch.zeros(
+ (1, 16, inputs_embeds.shape[1]),
+ dtype=torch.long,
+ device=inputs_embeds.device,
+ )
+ )
+
+ # # Remove EOS token if present
+ # if code[-1] == TALKER_CODEC_EOS_TOKEN_ID:
+ # code = code[:-1]
+
+ # Generate audio from codec codes
+ audio_tensors = []
+ for code in codes:
+ audio_tensor = self.generate_audio(code, voice_type)
+ audio_tensors.append(audio_tensor)
+ if len(audio_tensors) > 1:
+ logger.warning(
+ "Batched input for code2wav is not supported yet, only the first audio tensor will be returned"
+ )
+
+ return OmniOutput(
+ text_hidden_states=None,
+ multimodal_outputs={"model_outputs": audio_tensors[0].reshape(1, -1)},
+ )
+
+ # Fallback (shouldn't reach here)
+ return OmniOutput(
+ text_hidden_states=torch.zeros(
+ [inputs_embeds.shape[0], self.talker.config.hidden_size],
+ dtype=torch.bfloat16,
+ ).to(self._module_device(self.model)),
+ multimodal_outputs=None,
+ )
+
+ # ==================== Audio Generation ====================
+
+ def generate_audio(self, code: torch.Tensor, voice_type: str) -> torch.Tensor:
+ """
+ Generate audio waveform from codec codes.
+
+ Args:
+ code: [8, T] - 8-layer RVQ codec codes
+ voice_type: Voice type (not used in Qwen3, kept for compatibility)
+
+ Returns:
+ audio_tensor: [1, waveform_len] - Audio waveform
+ """
+ code2wav_dev = self._module_device(self.code2wav)
+
+ # Convert to tensor if needed
+ if isinstance(code, torch.Tensor):
+ talker_codes = code.to(dtype=torch.long, device=code2wav_dev)
+ else:
+ talker_codes = torch.as_tensor(code, dtype=torch.long, device=code2wav_dev)
+
+ # Ensure shape is [batch=1, 8, T]
+ if talker_codes.ndim == 2:
+ # [8, T] → [1, 8, T]
+ talker_codes = talker_codes.unsqueeze(0)
+ elif talker_codes.ndim == 1:
+ # [T] → assume single layer, expand to 16 layers
+ talker_codes = talker_codes.unsqueeze(0).unsqueeze(0)
+ talker_codes = talker_codes.expand(1, 16, -1)
+
+ # Use chunked decode for memory efficiency
+ audio_tensor = self.code2wav.chunked_decode(
+ talker_codes,
+ chunk_size=300,
+ left_context_size=25,
+ )
+
+ return audio_tensor
+
+ # ==================== Thinker-Talker Projection ====================
+
+ def _load_talker_embedding(self) -> torch.nn.Embedding:
+ """Load talker embedding layer."""
+ return self.talker.language_model.model.codec_embedding
+
+ def _init_special_tokens_embeddings(self) -> set[str]:
+ """
+ Initialize special token embeddings for thinker-talker projection.
+
+ Following Transformers implementation:
+ - TTS tokens (BOS/EOS/PAD) come from thinker's embedding, projected to talker space
+ - Codec tokens (BOS/EOS/PAD/NOTHINK/THINK_*) come from talker's embedding
+ - Speaker tokens are also from talker's embedding
+
+ Note on projections:
+ - text_projection: Used here for text token embeddings (thinker → talker dimension)
+ - hidden_projection: Used at runtime for multimodal hidden states (audio/image/video)
+ from thinker's last layer, not needed for special token initialization
+ """
+ # Get embeddings from both models
+ # self.thinker_embedding = self.thinker.model.get_input_embeddings()
+ self.talker_embedding = self._load_talker_embedding()
+
+ # Get configuration
+ talker_hf_config = self.talker_config
+ if hasattr(talker_hf_config, "talker_config"):
+ talker_hf_config = talker_hf_config.talker_config
+
+ codec_special_tokens = torch.tensor(
+ [
+ [
+ talker_hf_config.codec_nothink_id,
+ talker_hf_config.codec_think_bos_id,
+ talker_hf_config.codec_think_eos_id,
+ talker_hf_config.codec_pad_id,
+ talker_hf_config.codec_bos_id,
+ talker_hf_config.codec_eos_token_id,
+ ]
+ ],
+ device=self._module_device(self.talker),
+ dtype=torch.long,
+ )
+ codec_embeds = self.talker_embedding(codec_special_tokens) # [1, 6, talker_hidden]
+ (
+ self.embed_codec_nothink_token,
+ self.embed_codec_think_bos_token,
+ self.embed_codec_think_eos_token,
+ self.embed_codec_pad_token,
+ self.embed_codec_bos_token,
+ self.embed_codec_eos_token,
+ ) = codec_embeds.chunk(6, dim=1)
+
+ # Speaker token IDs (for voice selection)
+ # In Qwen3, speaker_id mapping is in talker_config.speaker_id
+ if hasattr(talker_hf_config, "speaker_id") and talker_hf_config.speaker_id:
+ self.tts_text_spk_token_ids = talker_hf_config.speaker_id
+ else:
+ # Default to audio_start_token_id if no speaker mapping
+ self.tts_text_spk_token_ids = {
+ "default": talker_hf_config.audio_start_token_id,
+ "Ethan": talker_hf_config.audio_start_token_id,
+ "prefix_caching": talker_hf_config.audio_start_token_id,
+ }
+
+ self.default_tts_text_spk_type = list(self.tts_text_spk_token_ids.keys())[0]
+
+ return set(["thinker_embedding.weight", "talker_embedding.weight"])
+
+ def _get_text_spk_token_id(self, voice_type: str) -> int:
+ """Get speaker token ID for voice type."""
+ if voice_type not in self.tts_text_spk_token_ids:
+ return self.tts_text_spk_token_ids[self.default_tts_text_spk_type]
+ return self.tts_text_spk_token_ids[voice_type]
+
+ def _thinker_to_talker_prefill(
+ self,
+ thinker_embed: torch.Tensor,
+ thinker_hidden: torch.Tensor,
+ multimodal_mask: Optional[torch.Tensor],
+ input_ids: torch.Tensor,
+ thinker_result_ids: torch.Tensor,
+ speaker_id,
+ tts_bos_thinker: Optional[torch.Tensor] = None,
+ tts_eos_thinker: Optional[torch.Tensor] = None,
+ tts_pad_thinker: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+ """
+ Project thinker outputs to talker inputs during prefill stage.
+
+ Returns:
+ (input_ids, input_embeds) for talker
+ """
+ im_start_indexes = torch.cat(
+ (
+ torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
+ torch.tensor([thinker_result_ids.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
+ ),
+ dim=-1,
+ ) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
+ multimodal_mask = (
+ (thinker_result_ids == self.thinker_config.audio_token_id) |
+ (thinker_result_ids == self.thinker_config.image_token_id) |
+ (thinker_result_ids == self.thinker_config.video_token_id)
+ ).to(input_ids.device) # [t] # fmt: skip
+
+ # Project thinker-side TTS embeddings into talker text space (provided by thinker stage)
+ talker_dev = self._module_device(self.talker)
+
+ def _ensure_1x1(x: torch.Tensor) -> torch.Tensor:
+ if x.ndim == 3:
+ return x[0, -1:, :]
+ if x.ndim == 2:
+ return x[-1]
+ return x.view(1, 1, -1)
+
+ def _proj_from_thinker(x_opt: Optional[torch.Tensor]) -> torch.Tensor:
+ if isinstance(x_opt, torch.Tensor) and x_opt.numel() > 0:
+ xin = _ensure_1x1(x_opt).to(talker_dev)
+ else:
+ xin = torch.zeros(
+ (1, thinker_embed.shape[-1]),
+ device=talker_dev,
+ dtype=thinker_embed.dtype,
+ )
+ return self.talker.text_projection(xin).to(input_ids.device)
+
+ tts_bos_embed = _proj_from_thinker(tts_bos_thinker)
+ tts_eos_embed = _proj_from_thinker(tts_eos_thinker)
+ tts_pad_embed = _proj_from_thinker(tts_pad_thinker)
+ self.tts_pad_embed = tts_pad_embed
+
+ talker_input_embeds = [] # [1 t d]
+ talker_input_ids = []
+ trailing_text_hidden_all: Optional[torch.Tensor] = None
+ # For every chatml parts
+ for i in range(len(im_start_indexes) - 1):
+ im_start_index = im_start_indexes[i].item()
+ segment_end_index = im_start_indexes[i + 1].item()
+ role_token = input_ids[0][im_start_index + 1]
+ # Talker should ignore thinker system prompt
+ if (role_token == self.config.system_token_id).item():
+ continue
+ # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
+ elif (role_token == self.config.user_token_id).item():
+ talker_user_part = self._get_talker_user_parts(
+ im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
+ )
+ talker_input_embeds.append(talker_user_part)
+ talker_input_ids.append(thinker_result_ids[im_start_index:segment_end_index])
+ # Take assistant output (for now)
+ elif (role_token == self.config.assistant_token_id).item() and i == len(im_start_indexes) - 2:
+ talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
+ im_start_index,
+ segment_end_index,
+ speaker_id,
+ thinker_embed,
+ tts_pad_embed,
+ tts_bos_embed,
+ tts_eos_embed,
+ )
+ talker_input_embeds.append(talker_assistant_embeds)
+ talker_input_ids.append(talker_assistant_ids)
+ # capture trailing text hidden for decode steps
+ try:
+ if isinstance(trailing_text_hidden, torch.Tensor):
+ trailing_text_hidden_all = trailing_text_hidden
+ except Exception:
+ pass
+ # History assistant output (ignore for now)
+ elif (role_token == self.config.assistant_token_id).item() and i != len(im_start_indexes) - 2:
+ continue
+ else:
+ raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
+ talker_input_embed = torch.cat([embed.to(input_ids.device) for embed in talker_input_embeds], dim=0)
+ talker_input_id = torch.cat([embed.to(input_ids.device) for embed in talker_input_ids], dim=0)
+
+ return talker_input_id, talker_input_embed, trailing_text_hidden_all
+
+ def _get_talker_user_parts(self, im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed):
+ user_talker_part = torch.empty(
+ (segment_end_index - im_start_index, self.config.talker_config.text_config.hidden_size),
+ device=thinker_hidden.device,
+ dtype=torch.bfloat16,
+ )
+
+ user_mm_mask = multimodal_mask[im_start_index:segment_end_index]
+
+ # Multimodal data exists
+ if user_mm_mask.any():
+ user_thinker_hidden_mm = thinker_hidden[im_start_index:segment_end_index][user_mm_mask]
+ mm_hidden = self.talker.hidden_projection(user_thinker_hidden_mm).to(thinker_hidden.device)
+ user_talker_part[user_mm_mask] = mm_hidden
+ user_thinker_embed = thinker_embed[im_start_index:segment_end_index][~user_mm_mask]
+ user_text_hidden = self.talker.text_projection(user_thinker_embed).to(thinker_hidden.device)
+ user_talker_part[~user_mm_mask] = user_text_hidden
+ return user_talker_part
+
+ def _get_talker_assistant_parts(
+ self, im_start_index, segment_end_index, speaker_id, thinker_embed, tts_pad_embed, tts_bos_embed, tts_eos_embed
+ ):
+ assistant_hidden = self.talker.text_projection(thinker_embed[im_start_index:segment_end_index]).to(
+ tts_pad_embed.device
+ ) # [t, d]
+ assistant_text_hidden = torch.cat(
+ (
+ assistant_hidden[:3],
+ tts_pad_embed.expand(4, -1),
+ tts_bos_embed,
+ assistant_hidden[3:4], # First text
+ ),
+ dim=0,
+ )
+ codec_special_tokens = torch.tensor(
+ [
+ self.config.talker_config.codec_nothink_id,
+ self.config.talker_config.codec_think_bos_id,
+ self.config.talker_config.codec_think_eos_id,
+ speaker_id,
+ self.config.talker_config.codec_pad_id,
+ self.config.talker_config.codec_bos_id,
+ ],
+ device=tts_pad_embed.device,
+ dtype=torch.long,
+ )
+ assistant_codec_hidden = torch.cat(
+ (
+ torch.zeros(
+ (3, self.config.talker_config.text_config.hidden_size),
+ device=tts_pad_embed.device,
+ dtype=torch.bfloat16,
+ ),
+ self.talker.get_input_embeddings(codec_special_tokens).to(
+ device=tts_pad_embed.device, dtype=torch.bfloat16
+ ),
+ ),
+ dim=0,
+ )
+ trailing_text_hidden = torch.cat(
+ (
+ assistant_hidden[4:],
+ tts_eos_embed,
+ ),
+ dim=0,
+ )
+
+ input_embeds = assistant_text_hidden + assistant_codec_hidden
+ input_ids = torch.full(
+ (assistant_text_hidden.shape[0],),
+ fill_value=self.config.tts_pad_token_id,
+ dtype=torch.long,
+ device=assistant_text_hidden.device,
+ )
+ return input_embeds, input_ids, trailing_text_hidden
+
+ def _talker_to_code_predictor(
+ self,
+ talker_hidden_states: Optional[torch.Tensor],
+ layer0_token_ids: Optional[torch.Tensor],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Project talker outputs to code predictor inputs.
+
+ Returns:
+ (input_ids, input_embeds) for code predictor.
+ """
+ predictor = getattr(self, "code_predictor", None)
+ device = (
+ self._module_device(predictor)
+ if predictor is not None
+ else (
+ talker_hidden_states.device
+ if isinstance(talker_hidden_states, torch.Tensor)
+ else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ )
+ )
+
+ if not isinstance(talker_hidden_states, torch.Tensor):
+ raise ValueError("Talker hidden states must be provided for the code predictor stage.")
+
+ inputs_embeds = talker_hidden_states.to(device=device, dtype=torch.bfloat16)
+ if inputs_embeds.ndim == 2:
+ inputs_embeds = inputs_embeds.unsqueeze(0)
+
+ if not isinstance(layer0_token_ids, torch.Tensor):
+ raise ValueError("Layer-0 codec token ids must accompany talker hidden states.")
+ input_ids = layer0_token_ids.to(device=device, dtype=torch.long)
+ if input_ids.ndim == 1:
+ input_ids = input_ids.unsqueeze(0)
+
+ return input_ids, inputs_embeds
+
+ # ==================== Logits and Sampling ====================
+
+ def compute_logits(
+ self,
+ hidden_states: Union[torch.Tensor, OmniOutput],
+ sampling_metadata: SamplingMetadata = None,
+ ) -> Optional[torch.Tensor]:
+ """Compute logits from hidden states."""
+ # Handle OmniOutput type
+ from vllm.v1.sample.logits_processor import LogitsProcessors
+
+ if isinstance(hidden_states, OmniOutput):
+ hidden_states = hidden_states.text_hidden_states
+ num_reqs = hidden_states.shape[0]
+ top_k = hidden_states.size(1) - 1
+
+ def dummy_tensors(v):
+ return torch.full((num_reqs,), v, device=hidden_states.device)
+
+ if sampling_metadata is None:
+ sampling_metadata = SamplingMetadata(
+ temperature=dummy_tensors(0.5),
+ all_greedy=False,
+ all_random=False,
+ top_p=dummy_tensors(0.9),
+ top_k=dummy_tensors(top_k),
+ generators={},
+ max_num_logprobs=None,
+ no_penalties=True,
+ prompt_token_ids=None,
+ frequency_penalties=dummy_tensors(0.1),
+ presence_penalties=dummy_tensors(0.1),
+ repetition_penalties=dummy_tensors(0.1),
+ output_token_ids=[[] for _ in range(num_reqs)],
+ allowed_token_ids_mask=None,
+ bad_words_token_ids={},
+ logitsprocs=LogitsProcessors(),
+ )
+ # Use active model for logits computation
+ logits = self.model.compute_logits(hidden_states, sampling_metadata) # V, d
+ # Talker: suppress tokens by setting their probability to ~1e-9 (finite very small),
+ # implemented by assigning their logits to log(1e-9).
+
+ if getattr(self, "model_stage", None) == "talker" and isinstance(logits, torch.Tensor):
+ # suppress tokens by setting their probability to ~1e-9 (finite very small)
+ suppressed_tokens = self._get_talker_suppressed_tokens()
+ try:
+ logits_cpu = logits.cpu()
+ logits_cpu[:, suppressed_tokens] = -1e9
+ logits = logits_cpu.to(logits.device)
+ except Exception as e:
+ print(f"Error in logits suppression: {e}")
+ print(f"logits.shape: {logits.shape}")
+ print(f"suppressed_tokens: {suppressed_tokens}")
+ raise e
+ logits[:, suppressed_tokens] = -1e9
+ return logits
+
+ def sample(
+ self,
+ logits: torch.Tensor,
+ sampling_metadata: SamplingMetadata,
+ ) -> Optional[SamplerOutput]:
+ """Sample from logits."""
+ return self.model.sample(logits, sampling_metadata)
+
+ # ==================== Weight Loading ====================
+
+ 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 = []
+ code2wav_weights = []
+
+ # Separate weights by component
+ 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("code2wav."):
+ code2wav_weights.append((k, v))
+ else:
+ logger.warning(f"Unknown weight prefix: {k}")
+ # Load thinker weights
+ if self.thinker and thinker_weights:
+ thinker_loaded = self.thinker.load_weights(thinker_weights)
+ thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker")
+ loaded_weights.update(thinker_loaded)
+
+ # Load talker weights
+ if self.talker and talker_weights:
+ 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 code2wav weights
+ if self.code2wav and code2wav_weights:
+ code2wav_loaded = self.code2wav.load_weights(code2wav_weights)
+ code2wav_loaded = add_prefix_to_loaded_weights(code2wav_loaded, "code2wav")
+ loaded_weights.update(code2wav_loaded)
+
+ # Log summary
+ logger.info(
+ "Loaded %d weights for Qwen3OmniMoe (stage=%s)",
+ len(loaded_weights),
+ self.model_stage,
+ )
+
+ return loaded_weights
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
new file mode 100644
index 000000000..869d50ae5
--- /dev/null
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright 2025 The Qwen team.
+"""Inference-only Qwen3-Omni-Moe Code2Wav model."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+ Qwen3OmniMoeCode2WavConfig,
+)
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+ Qwen3OmniMoeCausalConvNet,
+ Qwen3OmniMoeCausalTransConvNet,
+ Qwen3OmniMoeCode2WavDecoderBlock,
+ Qwen3OmniMoeCode2WavTransformerModel,
+ Qwen3OmniMoeConvNeXtBlock,
+ SnakeBeta,
+)
+from vllm.config import VllmConfig # type: ignore
+from vllm.logger import init_logger # type: ignore
+from vllm.model_executor.models.utils import ( # type: ignore
+ AutoWeightsLoader,
+ WeightsMapper,
+)
+
+logger = init_logger(__name__)
+
+
+class Qwen3OmniMoeCode2Wav(nn.Module):
+ """
+ Qwen3 Omni MoE Code2Wav - Converts num_quantizers-layer RVQ codec codes to audio waveform.
+
+ Architecture:
+ 1. Code Embedding: Embed and average num_quantizers RVQ layers
+ 2. Pre-Transformer: Add temporal context via sliding-window attention
+ 3. Upsampling: Progressive upsampling with ConvNeXt blocks
+ 4. Decoder: Multi-stage upsampling + residual units → waveform
+
+ Input: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes
+ Output: [batch, 1, waveform_len] - Audio waveform [-1, 1]
+
+ Total upsampling factor: ~1280x
+ Example: 100 codec frames → 128,000 audio samples (8 seconds at 16kHz)
+ """
+
+ input_modalities = "audio"
+
+ # Weight mapper
+ hf_to_vllm_mapper = WeightsMapper(
+ orig_to_new_prefix={
+ "code2wav.pre_transformer.": "pre_transformer.",
+ "code2wav.code_embedding.": "code_embedding.",
+ "code2wav.upsample.": "upsample.",
+ "code2wav.decoder.": "decoder.",
+ "code2wav.": "",
+ }
+ )
+
+ def __init__(
+ self,
+ *,
+ vllm_config: VllmConfig | None = None,
+ prefix: str = "",
+ ):
+ super().__init__()
+
+ self.config: Qwen3OmniMoeCode2WavConfig = vllm_config.model_config.hf_config
+
+ # Calculate total upsampling factor
+ self.total_upsample = np.prod(self.config.upsample_rates + self.config.upsampling_ratios)
+
+ # Pre-transformer
+ self.pre_transformer = Qwen3OmniMoeCode2WavTransformerModel._from_config(self.config)
+
+ # Code embedding: Single embedding table for all RVQ layers
+ self.code_embedding = nn.Embedding(
+ self.config.codebook_size * self.config.num_quantizers, self.config.hidden_size
+ )
+
+ # Offset for each RVQ layer (layer 0: 0-1023, layer 1: 1024-2047, etc.)
+ self.register_buffer(
+ "code_offset",
+ torch.arange(self.config.num_quantizers).view(1, -1, 1) * self.config.codebook_size,
+ persistent=False,
+ )
+
+ # Upsampling blocks (e.g., 2x, 2x)
+ upsample = []
+ for factor in self.config.upsampling_ratios:
+ upsample.append(
+ nn.ModuleList(
+ [
+ Qwen3OmniMoeCausalTransConvNet(
+ self.config.hidden_size, self.config.hidden_size, factor, factor
+ ),
+ Qwen3OmniMoeConvNeXtBlock(self.config.hidden_size),
+ ]
+ )
+ )
+ self.upsample = nn.ModuleList(upsample)
+
+ # Decoder: Initial projection + progressive upsampling blocks
+ decoder = [Qwen3OmniMoeCausalConvNet(self.config.hidden_size, self.config.decoder_dim, kernel_size=7)]
+
+ # Add decoder blocks (each upsamples and reduces channels)
+ for i in range(len(self.config.upsample_rates)):
+ decoder.append(Qwen3OmniMoeCode2WavDecoderBlock(self.config, i))
+
+ # Final projection to waveform
+ output_dim = self.config.decoder_dim // 2 ** len(self.config.upsample_rates)
+ decoder += [
+ SnakeBeta(output_dim),
+ Qwen3OmniMoeCausalConvNet(output_dim, 1, kernel_size=7),
+ ]
+ self.decoder = nn.ModuleList(decoder)
+
+ def forward(self, codes: torch.Tensor) -> torch.Tensor:
+ """
+ Convert num_quantizers-layer RVQ codes to audio waveform.
+
+ Args:
+ codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codec codes
+
+ Returns:
+ waveform: [batch, 1, waveform_len] - Audio waveform clipped to [-1, 1]
+ """
+ if codes.shape[1] != self.config.num_quantizers:
+ raise ValueError(f"Expected {self.config.num_quantizers} layers of codes, got {codes.shape[1]}")
+
+ # Stage 1: Code Embedding
+ # Add offset to separate layer vocabularies, then embed and average
+ hidden = self.code_embedding(codes + self.code_offset).mean(1)
+ # Shape: [batch, seq_len, hidden_size]
+
+ # Stage 2: Pre-Transformer (add temporal context)
+ hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state
+ # Shape: [batch, seq_len, hidden_size]
+
+ # Stage 3: Upsampling
+ hidden = hidden.permute(0, 2, 1) # [batch, hidden_size, seq_len]
+ for blocks in self.upsample:
+ for block in blocks:
+ hidden = block(hidden)
+ # Shape: [batch, hidden_size, seq_len * upsample_factor]
+
+ # Stage 4: Decoder (progressive upsampling to waveform)
+ wav = hidden
+ for block in self.decoder:
+ wav = block(wav)
+ # Shape: [batch, 1, waveform_len]
+
+ # Clamp to valid audio range
+ return wav.clamp(min=-1.0, max=1.0)
+
+ def chunked_decode(
+ self,
+ codes: torch.Tensor,
+ chunk_size: int = 300,
+ left_context_size: int = 25,
+ ) -> torch.Tensor:
+ """
+ Decode long sequences in chunks to avoid OOM.
+
+ Uses overlapping chunks with left context to avoid boundary artifacts.
+
+ Args:
+ codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes
+ chunk_size: Number of codec frames per chunk
+ left_context_size: Number of overlapping frames for context
+
+ Returns:
+ waveform: [batch, 1, waveform_len] - Complete waveform
+ """
+ wavs = []
+ start_index = 0
+
+ while start_index < codes.shape[-1]:
+ end_index = min(start_index + chunk_size, codes.shape[-1])
+ context_size = left_context_size if start_index >= left_context_size else start_index
+
+ # Extract chunk with left context
+ codes_chunk = codes[..., start_index - context_size : end_index]
+
+ # Decode chunk
+ wav_chunk = self(codes_chunk)
+
+ # Remove context from output (context_size * total_upsample samples)
+ wavs.append(wav_chunk[..., context_size * self.total_upsample :])
+
+ start_index = end_index
+
+ return torch.cat(wavs, dim=-1)
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ """Load weights from HuggingFace checkpoint."""
+ loader = AutoWeightsLoader(
+ self,
+ skip_prefixes=["thinker.", "talker."], # Already loaded above
+ )
+ 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
+ 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
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py
new file mode 100644
index 000000000..49215be36
--- /dev/null
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py
@@ -0,0 +1,491 @@
+"""Qwen3-Omni Code Predictor with MTP (Multi-Token Prediction) support.
+
+This module implements the code predictor component for Qwen3-Omni talker models.
+
+The code predictor generates residual RVQ (Residual Vector Quantization) codes
+autoregressively, predicting layers 1 to N based on layer-0 codes from the talker.
+"""
+
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import Cache, PretrainedConfig
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeRotaryEmbedding
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+
+logger = init_logger(__name__)
+
+
+# ============================================================================
+# Rotary Embeddings and Helper Functions
+# ============================================================================
+
+
+class Qwen3OmniCodePredictorRotaryEmbedding(nn.Module):
+ """Rotary positional embeddings for the code predictor."""
+
+ def __init__(self, dim, max_position_embeddings=2048, base=10000):
+ super().__init__()
+ self.dim = dim
+ self.max_position_embeddings = max_position_embeddings
+ self.base = base
+
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, x: torch.Tensor, position_ids: torch.LongTensor) -> tuple[torch.Tensor, torch.Tensor]:
+ inv_freq = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+ position = position_ids[:, None, :].float()
+ freqs = (inv_freq @ position).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos().to(x.dtype)
+ sin = emb.sin().to(x.dtype)
+ return cos, sin
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ position_ids (`torch.Tensor`, *optional*):
+ Deprecated and unused.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """Repeat key/value heads for grouped query attention."""
+ if n_rep == 1:
+ return hidden_states
+ bsz, num_kv_heads, seq_len, head_dim = hidden_states.shape
+ hidden_states = hidden_states[:, :, None, :, :].expand(bsz, num_kv_heads, n_rep, seq_len, head_dim)
+ return hidden_states.reshape(bsz, num_kv_heads * n_rep, seq_len, head_dim)
+
+
+# ============================================================================
+# Code Predictor Attention Layer
+# ============================================================================
+
+
+class Qwen3OmniCodePredictorAttention(nn.Module):
+ """Multi-head self-attention for code predictor with vLLM optimization."""
+
+ def __init__(self, config, layer_idx: int, vllm_config: VllmConfig = None):
+ super().__init__()
+
+ self.num_heads = config.code_predictor_config.num_attention_heads
+ self.num_key_value_heads = config.code_predictor_config.num_key_value_heads
+ self.head_dim = getattr(
+ config.code_predictor_config,
+ "head_dim",
+ config.code_predictor_config.hidden_size // config.code_predictor_config.num_attention_heads,
+ )
+ self.hidden_size = config.code_predictor_config.hidden_size
+
+ if self.num_heads % self.num_key_value_heads != 0:
+ raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+
+ # Projection layers
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+ # Query/Key normalization
+ self.q_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps)
+ self.k_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps)
+ self.is_causal = True
+ self.config = config
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ causal_mask: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ use_cache: bool = False,
+ position_ids: Optional[torch.LongTensor] = None,
+ ) -> torch.Tensor:
+ bsz, seq_len, _ = hidden_states.shape
+
+ # Project to Q, K, V
+ q = self.q_proj(hidden_states).view(bsz, seq_len, self.num_heads, self.head_dim)
+ k = self.k_proj(hidden_states).view(bsz, seq_len, self.num_key_value_heads, self.head_dim)
+ v = self.v_proj(hidden_states).view(bsz, seq_len, self.num_key_value_heads, self.head_dim)
+
+ # Apply normalization
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+
+ # Apply RoPE
+ q_heads = q.transpose(1, 2).contiguous()
+ k_heads = k.transpose(1, 2).contiguous()
+ q_heads, k_heads = apply_rotary_pos_emb(q_heads, k_heads, cos, sin)
+ v_heads = v.transpose(1, 2).contiguous()
+
+ if past_key_values is not None:
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ k_heads, v_heads = past_key_values.update(k_heads, v_heads, self.layer_idx, cache_kwargs)
+
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+ attention_interface = ALL_ATTENTION_FUNCTIONS["flash_attention_2"]
+ attn_output, _ = attention_interface(
+ self,
+ q_heads,
+ k_heads,
+ v_heads,
+ None,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.head_dim**-0.5,
+ sliding_window=None,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ output_hidden_states=True,
+ output_attentions=False,
+ )
+ attn_output = attn_output.reshape(*(hidden_states.shape[:-1]), -1).contiguous()
+
+ attn_output = self.o_proj(attn_output)
+ return attn_output
+
+
+# ============================================================================
+# Code Predictor MLP Layer
+# ============================================================================
+
+
+class Qwen3OmniCodePredictorMLP(nn.Module):
+ """Feed-forward network for code predictor."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.gate_proj = nn.Linear(
+ config.code_predictor_config.hidden_size, config.code_predictor_config.intermediate_size, bias=False
+ )
+ self.up_proj = nn.Linear(
+ config.code_predictor_config.hidden_size, config.code_predictor_config.intermediate_size, bias=False
+ )
+ self.down_proj = nn.Linear(
+ config.code_predictor_config.intermediate_size, config.code_predictor_config.hidden_size, bias=False
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ gate = F.silu(self.gate_proj(hidden_states))
+ up = self.up_proj(hidden_states)
+ return self.down_proj(gate * up)
+
+
+# ============================================================================
+# MTP Layer (Multi-Token Prediction Layer)
+# ============================================================================
+
+
+class Qwen3OmniCodePredictorMTPLayer(nn.Module):
+ """MTP layer for speculative decoding - predicts next residual code layer."""
+
+ def __init__(
+ self,
+ config: PretrainedConfig,
+ prefix: str,
+ model_config: ModelConfig,
+ layer_idx: int,
+ cache_config: Optional[CacheConfig] = None,
+ quant_config: Optional[QuantizationConfig] = None,
+ ) -> None:
+ super().__init__()
+ self.layer_idx = layer_idx
+ self.config = config
+
+ # Qwen3OmniCodePredictorDecoderLayer
+ self.self_attn = Qwen3OmniCodePredictorAttention(
+ config,
+ layer_idx,
+ vllm_config=type(
+ "VllmConfig",
+ (),
+ {"cache_config": cache_config, "quant_config": quant_config, "model_config": model_config},
+ )(),
+ )
+ self.mlp = Qwen3OmniCodePredictorMLP(config)
+ self.input_layernorm = RMSNorm(
+ config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps
+ )
+ self.post_attention_layernorm = RMSNorm(
+ config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps
+ )
+
+ def mtp_block(
+ self,
+ hidden_states: torch.Tensor,
+ causal_mask: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ past_key_values: Optional[Cache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ use_cache: bool = False,
+ position_ids: Optional[torch.LongTensor] = None,
+ ) -> torch.Tensor:
+ # Self-attention with residual
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ hidden_states = self.self_attn(
+ hidden_states, causal_mask, cos, sin, past_key_values, cache_position, use_cache, position_ids
+ )
+ hidden_states = residual + hidden_states
+
+ # MLP with residual
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+ def forward(
+ self,
+ input_ids: torch.Tensor,
+ positions: torch.Tensor,
+ previous_hidden_states: torch.Tensor,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ spec_step_index: int = 0,
+ ) -> torch.Tensor:
+ assert inputs_embeds is not None, "inputs_embeds required for MTP"
+
+ # Mask position 0 (not needed for MTP)
+ inputs_embeds[positions == 0] = 0
+
+ hidden_states = torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
+
+ # Get position info for RoPE
+ batch_size, seq_len, _ = hidden_states.shape
+ device = hidden_states.device
+ position_ids = torch.arange(seq_len, dtype=torch.long, device=device).unsqueeze(0).expand(batch_size, -1)
+
+ # Get RoPE embeddings
+ head_dim = self.self_attn.head_dim
+ rotary_emb = Qwen3OmniCodePredictorRotaryEmbedding(
+ head_dim, max_position_embeddings=self.config.code_predictor_config.max_position_embeddings
+ )
+ rotary_emb = Qwen3OmniMoeRotaryEmbedding(self.config)
+ cos, sin = rotary_emb(hidden_states, position_ids)
+
+ # Create causal mask
+ causal_mask = (
+ torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
+ .unsqueeze(0)
+ .unsqueeze(0)
+ )
+ causal_mask = causal_mask.masked_fill(causal_mask, float("-inf"))
+
+ # Forward through MTP block
+ hidden_states = self.mtp_block(hidden_states, causal_mask, cos, sin)
+
+ return hidden_states
+
+
+class Qwen3OmniCodePredictorBaseModel(nn.Module):
+ """
+ Base model for code predictor - matches HF Qwen3OmniMoeTalkerCodePredictorModel structure.
+
+ This is a simple transformer that processes inputs_embeds and outputs hidden states.
+ """
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ super().__init__()
+ config = vllm_config.model_config.hf_config.code_predictor_config
+
+ self.config = config
+ self.vocab_size = config.vocab_size
+ self.num_code_groups = config.num_code_groups
+
+ # Codec embeddings (for layers 1-num_code_groups-1)
+ self.codec_embedding = nn.ModuleList(
+ [
+ VocabParallelEmbedding(
+ config.vocab_size,
+ config.hidden_size,
+ )
+ for _ in range(config.num_code_groups - 1)
+ ]
+ )
+
+ # Decoder layers
+ self.layers = nn.ModuleList(
+ [
+ Qwen3OmniCodePredictorMTPLayer(
+ vllm_config.model_config.hf_config,
+ f"{prefix}.layers.{idx}",
+ model_config=vllm_config.model_config,
+ layer_idx=idx,
+ cache_config=vllm_config.cache_config,
+ quant_config=vllm_config.quant_config,
+ )
+ for idx in range(config.num_hidden_layers)
+ ]
+ )
+
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ # RoPE
+ self.rotary_emb = Qwen3OmniMoeRotaryEmbedding(config=config)
+
+ def forward(
+ self,
+ inputs_embeds: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Any] = None,
+ use_cache: Optional[bool] = False,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs,
+ ):
+ """
+ Forward pass matching HF structure.
+
+ Args:
+ inputs_embeds: [batch, seq_len, hidden_size]
+
+ Returns:
+ Object with .last_hidden_state attribute
+ """
+ batch_size, seq_len, _ = inputs_embeds.shape
+
+ # Create positions tensor if not provided
+ # positions must be [num_tokens] or [batch_size, seq_len]
+ if position_ids is None:
+ if cache_position is not None:
+ position_ids = cache_position # [num_tokens]
+ else:
+ position_ids = torch.arange(seq_len, device=inputs_embeds.device).unsqueeze(0) # [1, seq_len]
+ else:
+ position_ids = position_ids.flatten() # Ensure [num_tokens]
+
+ # Extract cos/sin from rotary_emb cache
+ # The cos_sin_cache is [max_pos, rotary_dim * 2]
+ cos, sin = self.rotary_emb(inputs_embeds, position_ids)
+
+ # Create causal mask
+ device = inputs_embeds.device
+ causal_mask = (
+ torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
+ .unsqueeze(0)
+ .unsqueeze(0)
+ )
+ causal_mask = causal_mask.masked_fill(causal_mask, float("-inf"))
+
+ # Forward through decoder layers
+ hidden_states = inputs_embeds
+
+ for layer in self.layers:
+ hidden_states = layer.mtp_block(
+ hidden_states, causal_mask, cos, sin, past_key_values, cache_position, use_cache, position_ids
+ )
+
+ # Final norm
+ hidden_states = self.norm(hidden_states)
+
+ # Return in HF-compatible format
+ from collections import namedtuple
+
+ Output = namedtuple("Output", ["last_hidden_state", "past_key_values"])
+ return Output(last_hidden_state=hidden_states, past_key_values=None) # [batch, num_code_groups-1, hidden_size]
+
+ def get_input_embeddings(self):
+ """Return codec embeddings for HF compatibility."""
+ return self.codec_embedding
+
+
+class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
+ """
+ Code predictor wrapper matching HF structure.
+
+ Structure:
+ - self.model: Qwen3OmniCodePredictorBaseModel (transformer)
+ - self.lm_head: ModuleList of output heads
+ """
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ super().__init__()
+
+ talker_code_predictor_config = vllm_config.model_config.hf_config
+ self.quant_config = vllm_config.quant_config
+ self.prefix = prefix
+
+ self.config = talker_code_predictor_config
+ self.vocab_size = self.config.code_predictor_config.vocab_size
+ self.num_code_groups = self.config.code_predictor_config.num_code_groups
+
+ # Base transformer model (matches HF structure)
+ self.model = Qwen3OmniCodePredictorBaseModel(vllm_config=vllm_config, prefix=prefix)
+
+ # Output heads for each residual layer (1-num_layers-1)
+ self.lm_head = nn.ModuleList(
+ [
+ nn.Linear(
+ self.config.code_predictor_config.hidden_size,
+ self.config.code_predictor_config.vocab_size,
+ bias=False,
+ )
+ for _ in range(self.num_code_groups - 1)
+ ]
+ )
+
+ def forward(
+ self,
+ inputs_embeds: torch.Tensor, # [batch, seq_len, hidden_size]
+ layer_idx: int, # Which layer to predict (0-num_layers-2 for layers 1-num_layers-1)
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Forward pass for code predictor.
+
+ Args:
+ inputs_embeds: Input embeddings [batch_size, seq_len, hidden_size]
+ layer_idx: Which residual layer to predict (0-num_layers-2 for layers 1-num_layers-1)
+
+ Returns:
+ logits: Predicted logits [batch_size, seq_len, vocab_size]
+ hidden_states: Output hidden states [batch_size, seq_len, hidden_size]
+ """
+ # Pass through base model
+ hidden_states = self.model(inputs_embeds)
+
+ # Get logits from corresponding head
+ logits = self.lm_head[layer_idx](hidden_states)
+
+ return logits, hidden_states
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
new file mode 100644
index 000000000..34e05925f
--- /dev/null
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
@@ -0,0 +1,703 @@
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+ Qwen3OmniMoeTalkerConfig,
+)
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+ Qwen3OmniMoeAudioEncoder,
+)
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.models.interfaces import (
+ MultiModalEmbeddings,
+ SupportsMultiModal,
+ SupportsPP,
+)
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+ Qwen2_5OmniThinkerDummyInputsBuilder,
+)
+from vllm.model_executor.models.qwen3_moe import Qwen3MoeMLP
+from vllm.model_executor.models.utils import (
+ AutoWeightsLoader,
+ WeightsMapper,
+ maybe_prefix,
+ sequence_parallel_chunk,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp import (
+ Qwen3OmniMoeTalkerCodePredictor,
+)
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
+ Qwen3MoeLLMForCausalLM,
+ Qwen3Omni_VisionTransformer,
+ Qwen3OmniMoeConditionalGenerationMixin,
+ Qwen3OmniMoeThinkerMultiModalProcessor,
+ Qwen3OmniMoeThinkerProcessingInfo,
+)
+
+try:
+ import flash_attn
+except (ImportError, ModuleNotFoundError):
+ flash_attn = None
+
+from vllm_omni.model_executor.models.utils import safe_tensor_reshape
+
+logger = init_logger(__name__)
+
+Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+ Qwen3OmniMoeThinkerMultiModalProcessor,
+ info=Qwen3OmniMoeThinkerProcessingInfo,
+ dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder,
+)
+class Qwen3OmniMoeTalkerForConditionalGeneration(
+ nn.Module,
+ SupportsMultiModal,
+ SupportsPP,
+ Qwen3OmniMoeConditionalGenerationMixin,
+):
+ """
+ Qwen3 Omni MoE Talker - Converts text to audio codec codes.
+
+ The talker is the second stage of Qwen3 Omni MoE's TTS pipeline:
+ 1. Thinker: Generates text response + hidden states
+ 2. Talker: Converts those to 8-layer audio codec codes
+ 3. Code2Wav: Converts codes to waveform
+
+ ## Key Components:
+ - text_projection: Projects thinker text embeddings → talker dimension
+ - hidden_projection: Projects thinker hidden states → talker dimension
+ - language_model: Main MoE transformer (generates layer 0)
+ - codec_head: Projects to codec vocabulary (layer 0 logits)
+ - code_predictor: Small transformer for layers 1-num_layers-1
+ """
+
+ logger = init_logger(__name__)
+
+ # Weight mapping from HuggingFace to vLLM naming convention
+ hf_to_vllm_mapper = WeightsMapper(
+ orig_to_new_prefix={
+ # Main MoE transformer model
+ "talker.model.": "language_model.model.",
+ # Codec head remains separate (outputs audio codes, not text)
+ "talker.codec_head.": "codec_head.",
+ # Code predictor: Now matches HF structure exactly (has .model sub-module)
+ # e.g., "talker.code_predictor.model.codec_embedding.0" → "code_predictor.model.codec_embedding.0"
+ "talker.code_predictor.": "code_predictor.",
+ # Projection layers
+ "talker.text_projection.": "text_projection.",
+ "talker.hidden_projection.": "hidden_projection.",
+ # Fallback: strip talker prefix
+ "talker.": "",
+ }
+ )
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ super().__init__()
+ talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config
+ self.quant_config = vllm_config.quant_config
+ self.prefix = prefix
+
+ self.config = talker_config
+ self.vocab_size = talker_config.text_config.vocab_size
+ self.router_aux_loss_coef = talker_config.text_config.router_aux_loss_coef
+ self.num_experts = talker_config.text_config.num_experts
+ self.num_experts_per_tok = talker_config.text_config.num_experts_per_tok
+ # thinker projection components for talker
+ self.text_projection = Qwen3OmniMoeTalkerResizeMLP(self.config)
+ self.hidden_projection = Qwen3OmniMoeTalkerResizeMLP(self.config)
+ self.codec_head = nn.Linear(self.config.text_config.hidden_size, self.config.text_config.vocab_size, bias=False)
+
+ self.rope_deltas = None
+ self.spatial_merge_size = self.config.spatial_merge_size
+
+ self.vocab_size = self.config.code_predictor_config.vocab_size
+ self.num_code_groups = self.config.code_predictor_config.num_code_groups
+
+ self.language_model = Qwen3OmniMoeModel(
+ vllm_config=vllm_config,
+ talker_config=self.config,
+ prefix=maybe_prefix(prefix, "language_model"),
+ )
+
+ self.code_predictor = Qwen3OmniMoeTalkerCodePredictor(
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor")
+ )
+
+ def code_predictor_forward(
+ self,
+ input_ids: torch.Tensor,
+ inputs_embeds: torch.Tensor | None = None,
+ *,
+ temperature: float = 1.0,
+ top_k: int = 50, # Match transformers default
+ top_p: float = 0.8, # Match transformers default
+ generation_steps: int | None = None,
+ last_talker_hidden: torch.Tensor | None = None,
+ **_: object,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Generate full RVQ codec codes for the provided sequence.
+
+ The code predictor consumes the layer-0 codec codes produced by the talker
+ alongside the talker's hidden states, and autoregressively predicts the remaining
+ residual layers (to num_codec_groups).
+
+ Returns:
+ tuple containing:
+ - residual_codes: A tensor of shape [batch, num_code_groups, seq_len] containing
+ the complete set of codec codes
+ - summed_embeddings: A tensor of shape [batch, seq_len, hidden_size]
+ Sum of all layer embeddings at each position (like Transformers)
+ """
+ if input_ids is None:
+ raise ValueError("`input_ids` containing layer-0 codec codes must be provided.")
+ if inputs_embeds is None:
+ raise ValueError("`inputs_embeds` containing talker hidden states must be provided.")
+
+ if inputs_embeds.ndim == 2:
+ inputs_embeds = inputs_embeds.unsqueeze(0)
+ if input_ids.ndim == 1:
+ input_ids = input_ids.unsqueeze(0)
+
+ # Ensure the tensors are contiguous for the autoregressive sampling loop
+ inputs_embeds = inputs_embeds.contiguous()
+ input_ids = input_ids.contiguous()
+
+ # Generate full codec codes using MTP
+ # This will be the parallel prediction implementation
+ batch_size, seq_len = input_ids.shape
+
+ # For now, use sequential generation (TODO: implement parallel)
+ # Result will be [batch, num_code_groups, seq_len]
+ # - all_codes_per_position will collect [batch, num_code_groups, 1] for each position
+ all_codes_per_position = []
+ middle_hidden_states = [] # Collect hidden states for each position
+
+ # Generate residual layers for each position
+ for pos in range(seq_len):
+ layer0_code = input_ids[:, pos : pos + 1] # [batch, 1]
+
+ # Predict all residual layers (layers 1 to num_code_groups-1) autoregressively
+ pos_codes = [layer0_code] # Start with layer 0: [batch, 1]
+
+ # Initial input: [last_talker_hidden, layer0_embed]
+ layer0_embed = self.get_input_embeddings(layer0_code)
+ prev_embed = layer0_embed # Track previous layer embedding
+ try:
+ current_input = torch.cat([last_talker_hidden, prev_embed], dim=1) # [batch, 2, hidden_size]
+ except Exception as e:
+ print(f"Error in current_input: {e}")
+ print(f"last_talker_hidden shape: {last_talker_hidden.shape}")
+ print(f"prev_embed shape: {prev_embed.shape}")
+ raise e
+
+ for layer_idx in range(self.num_code_groups - 1):
+ # Input for this layer: [last_talker_hidden, prev_layer_embed]
+
+ # Forward through code_predictor model
+ outputs = self.code_predictor.model(
+ inputs_embeds=current_input,
+ attention_mask=None,
+ position_ids=None,
+ past_key_values=None,
+ use_cache=False,
+ cache_position=None,
+ )
+
+ hidden_state = outputs.last_hidden_state # [batch, 2, hidden_size]
+
+ # Use the corresponding lm_head for this layer
+ logits = self.code_predictor.lm_head[layer_idx](hidden_state[:, -1:, :]) # [batch, 1, vocab_size]
+ from transformers.generation.logits_process import (
+ LogitsProcessorList,
+ TopKLogitsWarper,
+ TopPLogitsWarper,
+ )
+
+ logits_processors = LogitsProcessorList(
+ [
+ TopKLogitsWarper(top_k=top_k),
+ TopPLogitsWarper(top_p=top_p),
+ ]
+ )
+ input_ids_for_logits_processors = torch.tensor([pos_codes[1:]]).to(logits.device, dtype=torch.long)
+ logits = logits_processors(input_ids_for_logits_processors, logits.squeeze(0)).unsqueeze(0)
+
+ # Sample from the filtered distribution
+ probs = F.softmax(logits, dim=-1)
+ code = torch.multinomial(probs.squeeze(1), num_samples=1) # [batch, 1]
+ pos_codes.append(code)
+
+ # Update prev_embed for next layer (if not last layer)
+ # layer_idx=0 predicts layer 1, embed with codec_embedding[1]
+ new_embed = self.code_predictor.model.codec_embedding[layer_idx](code) # [batch, 1, hidden_size]
+ current_input = torch.cat([current_input, new_embed], dim=1) # [batch, 3~n, hidden_size]
+
+ # Stack all layers for this position: [batch, num_code_groups, 1]
+ pos_all_layers = torch.stack(pos_codes, dim=1) # [batch, num_code_groups, 1]
+ all_codes_per_position.append(pos_all_layers)
+ middle_hidden_states.append(current_input[:, 2:-1, :])
+
+ # Concatenate across positions: [batch, num_code_groups, seq_len]
+ result_codes = torch.cat(all_codes_per_position, dim=2)
+
+ # Build summed embeddings for each position (like Transformers)
+ # This combines layer-0 embed, mid layers hidden states, and last layer embed
+ all_summed_embeddings = []
+
+ for pos in range(seq_len):
+ # Layer 0 embedding
+ layer0_code = result_codes[:, 0, pos : pos + 1] # [batch, 1]
+ layer0_embed = self.get_input_embeddings(layer0_code) # [batch, 1, hidden_size]
+
+ # mid layers hidden states (from CodePredictor)
+ mid_residual_hiddens = middle_hidden_states[pos] # [batch, num_code_groups-2, hidden_size]
+ mid_list = list(mid_residual_hiddens.split(1, dim=1))
+
+ # last layer embedding
+ last_layer_code = result_codes[:, -1, pos : pos + 1] # [batch, 1]
+ last_residual_hidden = self.code_predictor.model.codec_embedding[-1](last_layer_code)
+
+ # Concatenate all layers: [batch, num_code_groups, hidden_size]
+ pos_codec_hiddens = torch.cat(
+ [layer0_embed] + mid_list + [last_residual_hidden],
+ dim=1,
+ )
+
+ # Sum across layers: [batch, 1, hidden_size] (like Transformers)
+ pos_summed = pos_codec_hiddens.sum(dim=1, keepdim=True)
+ all_summed_embeddings.append(pos_summed)
+
+ # Concatenate across positions: [batch, seq_len, hidden_size]
+ summed_embeddings = torch.cat(all_summed_embeddings, dim=1)
+
+ return result_codes, summed_embeddings
+
+ def init_multi_modal(self, thinker_config):
+ """
+ Initialize multimodal components from the thinker.
+
+ Unlike Qwen2.5 Omni which creates audio_tower and visual encoders here,
+ Qwen3 Omni MoE has a cleaner separation: the thinker is the ONLY module
+ that processes raw multimodal inputs. The talker only handles text-to-audio
+ conversion using pre-processed embeddings from the thinker.
+
+ This method exists for API compatibility and stores the thinker config
+ for reference. The actual multimodal processing components (audio_tower,
+ visual) are ONLY in the thinker, not duplicated in the talker.
+
+ Args:
+ thinker_config: Configuration from the thinker model (for reference only)
+ """
+ self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config)
+ self.visual = Qwen3Omni_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"),
+ # attn_backend_override=attn_backend_override,
+ )
+
+ def project_thinker_outputs(
+ self,
+ thinker_embeds: torch.Tensor | None = None,
+ thinker_hidden_states: torch.Tensor | None = None,
+ is_multimodal_mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ """
+ Project thinker outputs to talker's hidden dimension.
+
+ The talker has a different hidden size than the thinker, so we need
+ to project the inputs appropriately:
+ - Text embeddings (from thinker's embedding layer) → text_projection
+ - Hidden states (from thinker's last layer, for multimodal) → hidden_projection
+
+ Args:
+ thinker_embeds: Text embeddings from thinker [batch, seq, thinker_hidden]
+ thinker_hidden_states: Hidden states from thinker's last layer [batch, seq, thinker_hidden]
+ is_multimodal_mask: Boolean mask indicating multimodal positions [batch, seq]
+
+ Returns:
+ projected_embeds: [batch, seq, talker_hidden]
+ """
+ if thinker_embeds is None and thinker_hidden_states is None:
+ raise ValueError("Either thinker_embeds or thinker_hidden_states must be provided")
+
+ # If only embeddings provided, project all as text
+ if thinker_hidden_states is None or is_multimodal_mask is None:
+ return self.text_projection(thinker_embeds)
+
+ # If only hidden states provided, project all as hidden
+ if thinker_embeds is None:
+ return self.hidden_projection(thinker_hidden_states)
+
+ # Mixed case: use mask to decide which projection
+ batch_size, seq_len, _ = thinker_embeds.shape
+ output = torch.empty(
+ (batch_size, seq_len, self.config.text_config.hidden_size),
+ device=thinker_embeds.device,
+ dtype=thinker_embeds.dtype,
+ )
+
+ # Project multimodal regions using hidden states
+ if is_multimodal_mask.any():
+ mm_hidden = thinker_hidden_states[is_multimodal_mask]
+ projected_mm = self.hidden_projection(mm_hidden)
+ output[is_multimodal_mask] = projected_mm
+
+ # Project text regions using embeddings
+ if (~is_multimodal_mask).any():
+ text_embeds = thinker_embeds[~is_multimodal_mask]
+ projected_text = self.text_projection(text_embeds)
+ output[~is_multimodal_mask] = projected_text
+
+ return output
+
+ def forward(
+ self,
+ batched_talker_inputs: dict,
+ intermediate_tensors: IntermediateTensors | None = None,
+ **kwargs: object,
+ ) -> torch.Tensor | IntermediateTensors:
+ """Forward pass through the talker model.
+
+ For inference, the talker receives inputs_embeds that should already be
+ projected to talker's hidden dimension. If receiving raw thinker outputs,
+ use project_thinker_outputs() first.
+ """
+ # If intermediate_tensors is provided (pipeline parallel),
+ # inputs_embeds should be None
+ batched_input_ids = []
+ batched_positions = []
+ batched_input_embeds = []
+ batched_code_predictor_codes = []
+ if intermediate_tensors is not None:
+ inputs_embeds = None
+ for rid in batched_talker_inputs:
+ input_ids = safe_tensor_reshape(batched_talker_inputs[rid]["input_ids"], (1, -1))
+ positions = batched_talker_inputs[rid]["positions"]
+ inputs_embeds = safe_tensor_reshape(
+ batched_talker_inputs[rid]["inputs_embeds"], (-1, self.config.text_config.hidden_size)
+ )
+ text_step = safe_tensor_reshape(batched_talker_inputs[rid]["text_step"], (1, -1))
+ last_talker_hidden = safe_tensor_reshape(
+ batched_talker_inputs[rid]["last_talker_hidden"], (1, 1, self.config.text_config.hidden_size)
+ )
+ # for profiling
+ if inputs_embeds.shape[-1] == 2048:
+ inputs_embeds = self.text_projection(inputs_embeds)
+ if inputs_embeds.shape[0] == 1:
+ code_predictor_codes, summed_embeddings = self.code_predictor_forward(
+ input_ids, inputs_embeds.clone(), last_talker_hidden=last_talker_hidden
+ )
+ inputs_embeds = summed_embeddings.clone()
+ else:
+ code_predictor_codes = torch.zeros((0, self.num_code_groups), dtype=torch.long)
+ batched_input_ids.append(input_ids.reshape(-1))
+ batched_positions.append(positions)
+ batched_input_embeds.append((inputs_embeds + text_step).reshape(-1, self.config.text_config.hidden_size))
+ batched_code_predictor_codes.append(code_predictor_codes.squeeze(-1).detach().to("cpu").contiguous())
+ try:
+ talker_input_ids = torch.cat(batched_input_ids, dim=0)
+ talker_positions = torch.cat(batched_positions, dim=0)
+ talker_input_embeds = torch.cat(batched_input_embeds, dim=0)
+ except Exception as e:
+ print(f"Error in talker_input_embeds: {e}")
+ print(f"talker_input_embeds shape: {[tensor.shape for tensor in batched_input_embeds]}")
+ print(f"talker_input_ids shape: {[tensor.shape for tensor in batched_input_ids]}")
+ print(f"talker_positions shape: {[tensor.shape for tensor in batched_positions]}")
+ raise e
+ talker_hidden_states, _ = self.language_model.model(
+ talker_input_ids,
+ talker_positions,
+ intermediate_tensors,
+ inputs_embeds=talker_input_embeds,
+ **kwargs,
+ )
+
+ # Pass talker hidden states to code predictor
+ # Returns: (residual_codes, summed_embeddings=sum of all layer embeddings at each position)
+
+ # Return both talker hidden states and code predictor results
+ # code_predictor_codes: [batch, num_code_groups, seq_len]
+ # summed_embeddings: [batch, seq_len, hidden_size]
+ # - Sum of all layer embeddings at each position (like Transformers)
+
+ return talker_hidden_states, batched_code_predictor_codes
+
+ def compute_logits(
+ self,
+ hidden_states: torch.Tensor,
+ sampling_metadata: SamplingMetadata = None,
+ ) -> torch.Tensor | None:
+ """Compute logits for audio codec codes (layer 0 of RVQ).
+
+ This projects the hidden states to the codec vocabulary space.
+ For full audio generation, layers except 0 would be predicted by
+ the code_predictor after sampling.
+ """
+ logits = self.codec_head(hidden_states)
+ return logits
+
+ def make_empty_intermediate_tensors(
+ self, batch_size: int, dtype: torch.dtype, device: torch.device
+ ) -> IntermediateTensors:
+ """Create empty intermediate tensors for pipeline parallelism."""
+ return self.language_model.make_empty_intermediate_tensors(batch_size, dtype, device)
+
+ 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 | None:
+ 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.
+ # TODO: do projection for all multimodel
+ for modality in mm_input_by_modality:
+ multimodal_input = mm_input_by_modality[modality]
+ if modality == "image":
+ image_embeddings = self._process_image_input(multimodal_input)
+ image_embeddings = self.hidden_projection(image_embeddings)
+ multimodal_embeddings += tuple(image_embeddings)
+ if modality == "video":
+ video_embeddings = self._process_video_input(multimodal_input)
+ video_video_embeddings_project = ()
+ for video_embed in video_embeddings:
+ proj = nn.Linear(8192, 2048).to(device=video_embed.device, dtype=torch.bfloat16)
+ video_embed = proj(video_embed)
+ video_embed_project = self.hidden_projection(video_embed)
+ video_video_embeddings_project += (video_embed_project,)
+ multimodal_embeddings += tuple(video_video_embeddings_project)
+ if modality == "audio":
+ audio_embeddings = self._process_audio_input(multimodal_input)
+ audio_embeddings = self.hidden_projection(audio_embeddings)
+ multimodal_embeddings += tuple(audio_embeddings)
+ return multimodal_embeddings
+
+ def get_input_embeddings(self, input_ids: torch.Tensor, multimodal_embeddings: MultiModalEmbeddings | None = None):
+ """Get the input embedding layer (for codec tokens)."""
+ return self.language_model.get_input_embeddings(input_ids, multimodal_embeddings)
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ """Load weights for the talker model.
+
+ The weight mapping translates from HuggingFace naming convention
+ to vLLM's internal structure. Code predictor weights are routed
+ to its custom loader for vocab extension support.
+ """
+ loader = AutoWeightsLoader(
+ self,
+ skip_prefixes=["thinker.", "code2wav."],
+ # "code_predictor."],
+ )
+ # Don't apply mapper again since we already did it
+ 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
+ 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
+
+
+class Qwen3OmniMoeTalkerResizeMLP(nn.Module):
+ """
+ MLP for projecting between thinker and talker hidden dimensions.
+
+ The thinker and talker have different hidden sizes:
+ - Thinker: config.thinker_hidden_size (e.g., 3584)
+ - Talker: config.text_config.hidden_size (e.g., 2048)
+
+ This MLP projects from thinker → talker dimension.
+ Two instances are used:
+ - text_projection: For text embeddings from thinker's embedding layer
+ - hidden_projection: For hidden states from thinker's last transformer layer
+ """
+
+ def __init__(self, config: Qwen3OmniMoeTalkerConfig):
+ super().__init__()
+ self.linear_fc1 = nn.Linear(config.thinker_hidden_size, config.text_config.intermediate_size, bias=True)
+ self.linear_fc2 = nn.Linear(config.text_config.intermediate_size, config.text_config.hidden_size, bias=True)
+ self.act_fn = _ACTIVATION_REGISTRY[config.text_config.hidden_act] # silu
+
+ def forward(self, hidden_state):
+ return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
+
+
+class Qwen3OmniMoeModel(Qwen3MoeLLMForCausalLM):
+ def __init__(self, vllm_config, talker_config, prefix):
+ talker_vllm_config = vllm_config.with_hf_config(
+ talker_config.text_config, architectures=["Qwen3MoeForCausalLM"]
+ )
+ talker_vllm_config.model_config.hf_text_config = talker_vllm_config.model_config.hf_config
+ super().__init__(
+ vllm_config=talker_vllm_config,
+ prefix=prefix,
+ )
+
+ self.config = talker_config
+
+ # Remove the inherited LM head so the talker only exposes codec outputs.
+ if hasattr(self, "lm_head"):
+ del self.lm_head
+
+ # Replace the base embed tokens with codec embedding (defined below).
+ if hasattr(self.model, "embed_tokens"):
+ del self.model.embed_tokens
+
+ # Codec embedding for RVQ code generation
+ self.model.codec_embedding = nn.Embedding(
+ talker_config.text_config.vocab_size, talker_config.text_config.hidden_size
+ )
+
+ # Add shared expert to each MoE layer and patch the forward method
+ layer_idx = 0
+ for layer in self.model.layers:
+ # add shared expert to Qwen3OmniMoeSparseMoeBlock layers
+ if hasattr(layer.mlp, "experts"): # Check if it's a SparseMoeBlock
+ # Shared expert is a regular gated MLP (SwiGLU)
+ layer.mlp.shared_expert = Qwen3MoeMLP(
+ hidden_size=self.config.text_config.hidden_size,
+ intermediate_size=self.config.text_config.shared_expert_intermediate_size,
+ hidden_act=self.config.text_config.hidden_act,
+ quant_config=talker_vllm_config.quant_config,
+ reduce_results=False, # Don't reduce since we'll add it manually
+ prefix=f"{prefix}.layers.{layer_idx}.mlp.shared_expert",
+ )
+
+ # Shared expert gate outputs a single scalar per token
+ layer.mlp.shared_expert_gate = ReplicatedLinear(
+ self.config.text_config.hidden_size,
+ 1, # Output single scalar per token
+ bias=False,
+ quant_config=None,
+ prefix=f"{prefix}.layers.{layer_idx}.mlp.shared_expert_gate",
+ )
+
+ # Store MoE config values for router computation
+ layer.mlp.top_k = self.config.text_config.num_experts_per_tok
+ layer.mlp.norm_topk_prob = self.config.text_config.norm_topk_prob
+ layer.mlp.num_experts = self.config.text_config.num_experts
+
+ # Monkey-patch the forward method to use shared expert
+ layer.mlp.forward = self._create_moe_forward_with_shared_expert(layer.mlp)
+
+ layer_idx += 1
+
+ def _create_moe_forward_with_shared_expert(self, moe_layer):
+ """Create a forward method that includes shared expert computation.
+
+ This matches the Transformers implementation where:
+ 1. Compute shared expert output (regular MLP)
+ 2. Gate it with sigmoid(shared_expert_gate(x))
+ 3. Apply softmax BEFORE top-k selection (matches Transformers router)
+ 4. Add to routed expert outputs
+ """
+
+ def forward_with_shared_expert(hidden_states: torch.Tensor, layer_idx: int = 0) -> torch.Tensor:
+ # Save original shape
+ orig_shape = hidden_states.shape
+ hidden_dim = hidden_states.shape[-1]
+ hidden_states = hidden_states.view(-1, hidden_dim)
+
+ # handle sequence parallel if needed
+ if hasattr(moe_layer, "is_sequence_parallel") and moe_layer.is_sequence_parallel:
+ hidden_states = sequence_parallel_chunk(hidden_states)
+
+ # Compute shared expert output
+ # The shared expert is a regular MLP, not a routed MoE
+ shared_output = None
+ if hasattr(moe_layer, "shared_expert") and moe_layer.shared_expert is not None:
+ # Forward through shared expert MLP
+ shared_output = moe_layer.shared_expert(hidden_states)
+
+ # Apply gating with sigmoid: sigmoid(gate(x)) * shared_expert(x)
+ if hasattr(moe_layer, "shared_expert_gate") and moe_layer.shared_expert_gate is not None:
+ gate_logits, _ = moe_layer.shared_expert_gate(hidden_states)
+ gate_values = F.sigmoid(gate_logits) # [batch, 1]
+ shared_output = gate_values * shared_output # Broadcasting: [batch, 1] * [batch, hidden]
+
+ # Compute experts results
+ # router_logits: (num_tokens, n_experts)
+ router_logits, _ = moe_layer.gate(hidden_states)
+ experts_output = moe_layer.experts(hidden_states=hidden_states, router_logits=router_logits)
+
+ # combine experts and shared expert results
+ if shared_output is not None:
+ final_hidden_states = experts_output + shared_output
+
+ # Handle sequence parallel if needed
+ if hasattr(moe_layer, "is_sequence_parallel") and moe_layer.is_sequence_parallel:
+ from vllm.distributed import tensor_model_parallel_all_gather
+
+ num_tokens = orig_shape[0] if len(orig_shape) > 1 else 1
+ final_hidden_states = tensor_model_parallel_all_gather(final_hidden_states, 0)
+ final_hidden_states = final_hidden_states[:num_tokens]
+ try:
+ final_hidden_states.view(orig_shape)
+ except Exception as e:
+ print(f"Error viewing final hidden states: {e}")
+ print(f"final_hidden_states.shape: {final_hidden_states.shape}")
+ print(f"orig_shape: {orig_shape}")
+ raise e
+ # Return with original shape
+ return final_hidden_states.view(orig_shape)
+
+ return forward_with_shared_expert
+
+ def get_input_embeddings(
+ self,
+ input_ids: torch.Tensor,
+ multimodal_embeddings: MultiModalEmbeddings | None = None,
+ generation_steps=None,
+ ) -> torch.Tensor:
+ return self.model.codec_embedding(input_ids)
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
new file mode 100644
index 000000000..573dc42cc
--- /dev/null
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
@@ -0,0 +1,1668 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright 2025 The Qwen team.
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+"""Inference-only Qwen3-Omni-Moe model (thinker part)."""
+
+import math
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from functools import partial
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from packaging.version import Version
+from transformers import PretrainedConfig
+from transformers import __version__ as TRANSFORMERS_VERSION
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+ Qwen3OmniMoeConfig,
+ Qwen3OmniMoeThinkerConfig,
+)
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+ Qwen3OmniMoeAudioEncoder,
+)
+from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import (
+ Qwen3OmniMoeProcessor,
+)
+from transformers.models.whisper import WhisperFeatureExtractor
+from vllm.attention.utils.fa_utils import get_flash_attn_version
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed import get_pp_group
+from vllm.logger import init_logger
+from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
+from vllm.model_executor.layers.linear import (
+ ColumnParallelLinear,
+ ReplicatedLinear,
+ RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+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,
+)
+from vllm.model_executor.models.qwen2_5_vl import (
+ Qwen2_5_VisionAttention,
+ Qwen2_5_VisionRotaryEmbedding,
+ Qwen2_5_VLProcessingInfo,
+)
+from vllm.model_executor.models.qwen2_audio import Qwen2AudioInputs, Qwen2AudioProcessingInfo
+from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM
+from vllm.model_executor.models.utils import (
+ AutoWeightsLoader,
+ WeightsMapper,
+ _merge_multimodal_embeddings,
+ flatten_bn,
+ maybe_prefix,
+)
+from vllm.model_executor.models.vision import (
+ get_vit_attn_backend,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import (
+ MultiModalKwargsItems,
+)
+from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems
+from vllm.multimodal.processing import (
+ MultiModalPromptUpdates,
+ PlaceholderFeaturesInfo,
+ PromptReplacement,
+ PromptUpdate,
+ PromptUpdateDetails,
+)
+from vllm.platforms import _Backend
+from vllm.sequence import IntermediateTensors
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from .qwen3_moe import Qwen3MoeModel
+
+try:
+ import flash_attn
+except (ImportError, ModuleNotFoundError):
+ flash_attn = None
+
+from torch import Tensor
+
+logger = init_logger(__name__)
+
+
+# Local helper functions (not available in vLLM)
+def conv3d_to_linear_weight(weight: torch.Tensor) -> torch.Tensor:
+ """Convert 3D convolutional weight to linear weight for patch embedding."""
+ # Weight shape: (out_channels, in_channels, t, h, w)
+ # Need to flatten spatial dimensions: (out_channels, in_channels * t * h * w)
+ flattened_weight = weight.flatten(1) # Flatten all dims except out_channels
+ return flattened_weight
+
+
+def get_llm_pos_ids_for_vision(
+ start_idx: int,
+ vision_idx: int,
+ spatial_merge_size: int,
+ t_index: list[torch.Tensor],
+ grid_hs: list[torch.Tensor],
+ grid_ws: list[torch.Tensor],
+) -> torch.Tensor:
+ """Generate 3D position IDs for vision tokens (temporal, height, width)."""
+ 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().float()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten().float()
+ t_index_tensor = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().float()
+ llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
+ llm_pos_ids = llm_pos_ids + start_idx
+ return llm_pos_ids
+
+
+# Type aliases for audio inputs
+Qwen2_5OmniAudioFeatureInputs = Qwen2AudioInputs
+
+
+def _get_feat_extract_output_lengths(input_lengths: torch.Tensor):
+ input_lengths_leave = input_lengths % 100
+ feat_lengths = (input_lengths_leave - 1) // 2 + 1
+ output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
+ return feat_lengths, output_lengths
+
+
+class Qwen3_VisionPatchEmbed(nn.Module):
+ def __init__(
+ self,
+ patch_size: int = 14,
+ temporal_patch_size: int = 2,
+ in_channels: int = 3,
+ hidden_size: int = 1152,
+ ) -> None:
+ super().__init__()
+ self.patch_size = patch_size
+ self.temporal_patch_size = temporal_patch_size
+ self.hidden_size = hidden_size
+
+ kernel_size = (temporal_patch_size, patch_size, patch_size)
+ self.proj = ReplicatedLinear(
+ in_channels * math.prod(kernel_size),
+ hidden_size,
+ bias=True,
+ return_bias=False,
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ L, C = x.shape
+ x = self.proj(x)
+ return x
+
+
+class Qwen3_VisionMLP(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ hidden_features: int,
+ bias: bool = False,
+ act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ):
+ super().__init__()
+ self.linear_fc1 = ColumnParallelLinear(
+ in_features,
+ hidden_features,
+ bias=bias,
+ quant_config=quant_config,
+ return_bias=False,
+ prefix=f"{prefix}.linear_fc1",
+ )
+ self.linear_fc2 = RowParallelLinear(
+ hidden_features,
+ in_features,
+ bias=bias,
+ quant_config=quant_config,
+ return_bias=False,
+ prefix=f"{prefix}.linear_fc2",
+ )
+ self.act_fn = act_fn
+
+ def forward(self, x: torch.Tensor):
+ mlp_output = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
+ return mlp_output
+
+
+class Qwen3_VisionBlock(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int,
+ mlp_hidden_dim: int,
+ act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,
+ norm_layer: Callable[[int], nn.Module] | None = None,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ if norm_layer is None:
+ norm_layer = partial(nn.LayerNorm, eps=1e-6)
+ self.norm1 = norm_layer(dim)
+ self.norm2 = norm_layer(dim)
+ self.attn = Qwen2_5_VisionAttention(
+ embed_dim=dim,
+ num_heads=num_heads,
+ projection_size=dim,
+ quant_config=quant_config,
+ prefix=f"{prefix}.attn",
+ )
+ self.mlp = Qwen3_VisionMLP(
+ dim,
+ mlp_hidden_dim,
+ act_fn=act_fn,
+ bias=True,
+ quant_config=quant_config,
+ prefix=f"{prefix}.mlp",
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: torch.Tensor,
+ max_seqlen: int | None = None, # Only used for Flash Attention
+ seqlens: list[int] | None = None, # Only used for xFormers
+ ) -> torch.Tensor:
+ x = x + self.attn(
+ self.norm1(x),
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ max_seqlen=max_seqlen,
+ seqlens=seqlens,
+ )
+
+ x = x + self.mlp(self.norm2(x))
+ return x
+
+
+class Qwen3_VisionPatchMerger(nn.Module):
+ def __init__(
+ self,
+ d_model: int,
+ context_dim: int,
+ norm_layer: Callable[[int], nn.Module] | None = None,
+ spatial_merge_size: int = 2,
+ use_postshuffle_norm: bool = False,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.hidden_size = context_dim * (spatial_merge_size**2)
+
+ self.use_postshuffle_norm = use_postshuffle_norm
+ if self.use_postshuffle_norm:
+ context_dim = self.hidden_size
+
+ if norm_layer is None:
+ norm_layer = partial(nn.LayerNorm, eps=1e-6)
+ self.use_postshuffle_norm = use_postshuffle_norm
+ self.ln_q = norm_layer(self.hidden_size if use_postshuffle_norm else context_dim)
+ self.mlp = nn.ModuleList(
+ [
+ ColumnParallelLinear(
+ self.hidden_size,
+ self.hidden_size,
+ bias=True,
+ quant_config=quant_config,
+ prefix=f"{prefix}.mlp.0",
+ ),
+ nn.GELU(),
+ RowParallelLinear(
+ self.hidden_size,
+ d_model,
+ bias=True,
+ quant_config=quant_config,
+ prefix=f"{prefix}.mlp.2",
+ ),
+ ]
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.use_postshuffle_norm:
+ x = self.ln_q(x.view(-1, self.hidden_size))
+ else:
+ x = self.ln_q(x).view(-1, self.hidden_size)
+
+ mlp_fc1, mlp_act, mlp_fc2 = self.mlp
+ x_parallel, _ = mlp_fc1(x)
+ x_parallel = mlp_act(x_parallel)
+ out, _ = mlp_fc2(x_parallel)
+ return out
+
+
+class Qwen3Omni_VisionTransformer(nn.Module):
+ def __init__(
+ self,
+ vision_config,
+ norm_eps: float = 1e-6,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ attn_backend_override: _Backend | None = None,
+ ) -> None:
+ super().__init__()
+ self.hidden_size = vision_config.hidden_size
+ self.num_heads = vision_config.num_heads
+ self.image_size = vision_config.image_size
+ self.patch_size = vision_config.patch_size
+ self.spatial_merge_size = vision_config.spatial_merge_size
+ self.spatial_merge_unit = self.spatial_merge_size**2
+ self.temporal_patch_size = vision_config.temporal_patch_size
+ self.num_grid_per_side = self.image_size // self.patch_size
+ self.apply_vit_abs_pos_embed = vision_config.apply_vit_abs_pos_embed
+ self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes
+
+ self.patch_embed = Qwen3_VisionPatchEmbed(
+ patch_size=self.patch_size,
+ temporal_patch_size=self.temporal_patch_size,
+ in_channels=vision_config.in_channels,
+ hidden_size=self.hidden_size,
+ )
+
+ # vit pos embedding, TODO: spatial_patch_size vs patch_size
+ if self.apply_vit_abs_pos_embed:
+ self.pos_embed = nn.Embedding(self.num_grid_per_side**2, self.hidden_size)
+ else:
+ self.pos_embed = nn.Parameter(torch.empty([1, self.num_grid_per_side**2, self.hidden_size]))
+
+ norm_layer = partial(nn.LayerNorm, eps=norm_eps)
+ head_dim = self.hidden_size // self.num_heads
+ self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
+
+ self.blocks = nn.ModuleList(
+ [
+ Qwen3_VisionBlock(
+ dim=self.hidden_size,
+ num_heads=self.num_heads,
+ mlp_hidden_dim=vision_config.intermediate_size,
+ act_fn=_ACTIVATION_REGISTRY[vision_config.hidden_act],
+ norm_layer=norm_layer,
+ quant_config=quant_config,
+ prefix=f"{prefix}.blocks.{layer_idx}",
+ )
+ for layer_idx in range(vision_config.depth)
+ ]
+ )
+ self.merger = Qwen3_VisionPatchMerger(
+ d_model=vision_config.out_hidden_size,
+ context_dim=self.hidden_size,
+ norm_layer=norm_layer,
+ spatial_merge_size=self.spatial_merge_size,
+ quant_config=quant_config,
+ prefix=f"{prefix}.merger",
+ )
+ if self.deepstack_visual_indexes is not None:
+ self.merger_list = nn.ModuleList(
+ [
+ Qwen3_VisionPatchMerger(
+ d_model=vision_config.out_hidden_size,
+ context_dim=self.hidden_size,
+ spatial_merge_size=self.spatial_merge_size,
+ use_postshuffle_norm=True,
+ norm_layer=norm_layer,
+ quant_config=quant_config,
+ prefix=f"{prefix}.merger_list.{layer_idx}",
+ )
+ for layer_idx in range(len(self.deepstack_visual_indexes))
+ ]
+ )
+
+ self.attn_backend = get_vit_attn_backend(
+ head_size=head_dim,
+ dtype=torch.get_default_dtype(),
+ # attn_backend_override=attn_backend_override,
+ )
+ if self.attn_backend != _Backend.FLASH_ATTN and get_flash_attn_version() is not None:
+ self.attn_backend = _Backend.FLASH_ATTN
+
+ @property
+ def dtype(self) -> torch.dtype:
+ return self.patch_embed.proj.weight.dtype
+
+ @property
+ def device(self) -> torch.device:
+ return self.patch_embed.proj.weight.device
+
+ def rot_pos_emb(self, grid_thw):
+ pos_ids = []
+ for t, h, w in grid_thw:
+ hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+ hpos_ids = hpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+ hpos_ids = hpos_ids.flatten()
+
+ wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+ wpos_ids = wpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+ wpos_ids = wpos_ids.flatten()
+ pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+ pos_ids = torch.cat(pos_ids, dim=0)
+ max_grid_size = grid_thw[:, 1:].max()
+ rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
+ rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
+ return rotary_pos_emb
+
+ def fast_pos_embed_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor:
+ num_grid_per_side = self.num_grid_per_side
+ m_size = self.spatial_merge_size
+ hidden_dim = self.pos_embed.embedding_dim
+
+ outputs = []
+ for t, h, w in grid_thw:
+ h_idxs = torch.linspace(0, num_grid_per_side - 1, h, dtype=torch.float32, device=self.device)
+ w_idxs = torch.linspace(0, num_grid_per_side - 1, w, dtype=torch.float32, device=self.device)
+
+ h_floor = h_idxs.to(torch.long)
+ w_floor = w_idxs.to(torch.long)
+ h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1)
+ w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1)
+
+ dh = h_idxs - h_floor
+ dw = w_idxs - w_floor
+
+ # Create meshgrid view for all h, w vars
+ dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij")
+ h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij")
+ h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij")
+ h_floor_grid_idx = h_floor_grid * num_grid_per_side
+ h_ceil_grid_idx = h_ceil_grid * num_grid_per_side
+
+ # original computation of weights
+ # w00 = (1 - dh_grid) * (1 - dw_grid)
+ # w01 = (1 - dh_grid) * dw_grid
+ # w10 = dh_grid * (1 - dw_grid)
+ # w11 = dh_grid * dw_grid
+ # we reuse w11 here to avoid duplicate
+ # dh_grid * dw_grid computation
+ w11 = dh_grid * dw_grid
+ w10 = dh_grid - w11
+ w01 = dw_grid - w11
+ w00 = 1 - dh_grid - dw_grid + w11
+
+ idx00 = h_floor_grid_idx + w_floor_grid
+ idx01 = h_floor_grid_idx + w_ceil_grid
+ idx10 = h_ceil_grid_idx + w_floor_grid
+ idx11 = h_ceil_grid_idx + w_ceil_grid
+
+ indices = torch.stack([idx00, idx01, idx10, idx11], dim=0).reshape(4, -1)
+ weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1)
+ weights = weights.to(dtype=self.dtype, device=self.device)
+
+ embeds = self.pos_embed(indices)
+ weighted_embeds = embeds * weights
+ p0, p1, p2, p3 = weighted_embeds.unbind(dim=0)
+ combined = p0 + p1 + p2 + p3
+
+ combined = combined.view(h * w, hidden_dim)
+ repeated = combined.unsqueeze(0).expand(t, -1, -1).contiguous()
+ repeated = repeated.view(t, h // m_size, m_size, w // m_size, m_size, hidden_dim)
+ repeated = repeated.permute(0, 1, 3, 2, 4, 5).reshape(-1, hidden_dim)
+ outputs.append(repeated)
+
+ return torch.cat(outputs, dim=0)
+
+ def compute_attn_mask_seqlen(
+ self,
+ cu_seqlens: torch.Tensor,
+ ) -> tuple[int | None, list[int] | None]:
+ max_seqlen, seqlens = None, None
+ if self.attn_backend == _Backend.FLASH_ATTN:
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+ elif self.attn_backend == _Backend.XFORMERS:
+ seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+ return max_seqlen, seqlens
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ grid_thw: list[list[int]],
+ ) -> torch.Tensor:
+ hidden_states = x.to(device=self.device, dtype=self.dtype)
+ hidden_states = self.patch_embed(hidden_states)
+
+ if self.apply_vit_abs_pos_embed:
+ pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
+ hidden_states = hidden_states + pos_embeds
+ rotary_pos_emb = self.rot_pos_emb(grid_thw)
+
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ dim=0,
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+ hidden_states = hidden_states.unsqueeze(1)
+ rotary_pos_emb = rotary_pos_emb.to(hidden_states.device)
+ max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
+
+ hidden_states_list = []
+ deepstack_visual_indexes = self.deepstack_visual_indexes
+
+ for layer_num, blk in enumerate(self.blocks):
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ max_seqlen=max_seqlen,
+ seqlens=seqlens,
+ )
+ if deepstack_visual_indexes is not None and layer_num in deepstack_visual_indexes:
+ hidden_states_list.append(hidden_states)
+
+ hidden_states = self.merger(hidden_states)
+
+ # processing deepstack
+ if deepstack_visual_indexes is not None:
+ processed_hidden_states_list = [hidden_states]
+ for idx, x in enumerate(hidden_states_list):
+ x = self.merger_list[idx](x)
+ processed_hidden_states_list.append(x)
+ # we cat the original visual features and deepstack features
+ # along the feature dim
+ hidden_states = torch.cat(
+ processed_hidden_states_list, dim=1
+ ) # [seq_len, hidden_size * (1 + depth_of_deepstack)]
+
+ return hidden_states
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ stacked_params_mapping = [
+ # (param_name, shard_name, shard_id)
+ ("attn.qkv.", "attn.q.", "q"),
+ ("attn.qkv.", "attn.k.", "k"),
+ ("attn.qkv.", "attn.v.", "v"),
+ ]
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
+ loaded_params: set[str] = set()
+
+ for name, loaded_weight in weights:
+ if name.endswith("patch_embed.proj.weight"):
+ loaded_weight = conv3d_to_linear_weight(loaded_weight)
+
+ 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)
+
+ param = params_dict[name]
+ weight_loader = param.weight_loader
+ weight_loader(param, loaded_weight, shard_id)
+ break
+ else:
+ 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
+
+
+@support_torch_compile(
+ dynamic_arg_dims={
+ "input_ids": 0,
+ "positions": -1,
+ "intermediate_tensors": 0,
+ "inputs_embeds": 0,
+ "deepstack_input_embeds": 0,
+ }
+)
+class Qwen3MoeLLMModel(Qwen3MoeModel):
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
+ # from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeTextRMSNorm
+ # config = vllm_config.model_config.hf_config
+ # self.norm = Qwen3OmniMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.deepstack_multiscale_layer_start = 1
+
+ def _deepstack_process(self, hidden_states: torch.Tensor, visual_embeds: torch.Tensor) -> torch.Tensor:
+ """
+ Process deepstack visual embeddings by selectively adding them only to visual positions.
+ This matches the transformers implementation which uses visual_pos_masks.
+
+ Args:
+ hidden_states: Hidden states tensor [num_tokens, hidden_size]
+ visual_embeds: Visual embeddings tensor [num_tokens, hidden_size] (zero-padded for non-visual positions)
+
+ Returns:
+ Modified hidden_states with visual embeddings added only to visual positions
+ """
+ # Derive visual position mask from visual_embeds (non-zero positions are visual)
+ # Use a small threshold to handle floating point precision issues
+ visual_pos_masks = torch.any(torch.abs(visual_embeds) > 1e-8, dim=-1)
+
+ if visual_pos_masks.any():
+ visual_pos_masks = visual_pos_masks.to(hidden_states.device)
+ visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
+ # Match transformers implementation EXACTLY (lines 2966-2969):
+ # In transformers: local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds
+ # Note: visual_embeds in transformers is NOT indexed, meaning it's already shaped
+ # [num_visual_tokens, hidden_size]. In vllm-omni, visual_embeds is [num_tokens, hidden_size]
+ # zero-padded, so we extract only visual positions to match transformers behavior.
+ visual_embeds_only = visual_embeds[visual_pos_masks, :] # Extract only visual positions
+ local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds_only
+ hidden_states[visual_pos_masks, :] = local_this
+
+ return hidden_states
+
+ def forward(
+ self,
+ input_ids: torch.Tensor,
+ positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ *,
+ capture_layer_indices: Sequence[int] | None = None,
+ return_hidden_states: bool = False,
+ deepstack_input_embeds: IntermediateTensors | None = None,
+ ) -> 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"]
+ capture_set = set(capture_layer_indices) if capture_layer_indices else None
+ captured_hidden_states: dict[str, torch.Tensor] | None = {} if return_hidden_states else None
+
+ for layer_idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
+ layer_idx = layer_idx + self.start_layer
+
+ if captured_hidden_states is not None and capture_set is not None:
+ if layer_idx in capture_set:
+ captured_hidden_states[str(layer_idx)] = hidden_states.clone().view(-1, hidden_states.shape[-1])
+
+ hidden_states, residual = layer(
+ positions,
+ hidden_states,
+ residual,
+ )
+
+ if deepstack_input_embeds is not None and layer_idx in range(0, len(deepstack_input_embeds)):
+ # Use _deepstack_process to match transformers implementation
+ # This selectively adds visual embeddings only to visual positions
+ hidden_states = self._deepstack_process(
+ hidden_states, deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]
+ )
+
+ if not get_pp_group().is_last_rank:
+ return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
+ hidden_states, _ = self.norm(hidden_states, residual)
+ if captured_hidden_states is not None:
+ return hidden_states, captured_hidden_states
+ else:
+ return hidden_states, None
+
+
+class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM):
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ super(Qwen3MoeForCausalLM, self).__init__()
+ config = vllm_config.model_config.hf_config
+ quant_config = vllm_config.quant_config
+ self.config = config
+ self.quant_config = quant_config
+ self.model = Qwen3MoeLLMModel(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+ self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, quant_config=quant_config)
+ if self.config.tie_word_embeddings:
+ self.lm_head.weight = self.model.embed_tokens.weight
+ self.logits_processor = LogitsProcessor(config.vocab_size)
+ self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
+
+
+class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo):
+ def get_hf_config(self):
+ return self.ctx.get_hf_config(Qwen3OmniMoeConfig).thinker_config
+
+ def get_hf_processor(self, **kwargs: object) -> Qwen3OmniMoeProcessor:
+ processor = self.ctx.get_hf_processor(
+ Qwen3OmniMoeProcessor,
+ use_fast=kwargs.pop("use_fast", True),
+ **kwargs,
+ )
+ if not hasattr(processor, "audio_token"):
+ processor.audio_token = "<|audio_pad|>"
+ if not hasattr(processor, "image_token"):
+ processor.image_token = "<|image_pad|>"
+ if not hasattr(processor, "video_token"):
+ processor.video_token = "<|video_pad|>"
+ return processor
+
+ 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, int | None]:
+ return {"audio": None, "image": None, "video": None}
+
+
+Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder
+
+
+class Qwen3OmniMoeThinkerMultiModalProcessor(
+ Qwen2_5OmniThinkerMultiModalProcessor,
+):
+ 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", [])
+
+ def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray:
+ length = x.shape[-1]
+ if length % hop_length != 0:
+ pad_length = hop_length - (length % hop_length)
+ x = np.pad(x, (0, pad_length), mode="constant", constant_values=0)
+ return x
+
+ # NOTE: WhisperFeatureExtractor cannot handle empty list of audios
+ feature_extractor = self.info.get_feature_extractor()
+ hop_length = feature_extractor.hop_length
+ if audios:
+ # NOTE: Qwen3-Omni processor accept "audio"
+ # To make sure the cache works with padding=True, we pre-padded
+ # the audio to multiple of hop_length.
+ mm_data["audio"] = [
+ pad_to_hop_length(audio, hop_length)
+ if isinstance(audio, np.ndarray)
+ else (pad_to_hop_length(audio[0], hop_length), audio[1])
+ for audio in audios
+ ]
+
+ # TODO(Isotr0py): Remove this patch after upstream fix PR
+ # released and Transformers version update:
+ # https://github.com/huggingface/transformers/pull/41473
+ mm_kwargs = dict(mm_kwargs)
+ tok_kwargs = dict(tok_kwargs)
+ if Version(TRANSFORMERS_VERSION) < Version("4.58.0"):
+ # move truncation to audio_kwargs level to avoid conflict
+ # with tok_kwargs
+ mm_kwargs["audio_kwargs"] = {"truncation": mm_kwargs.pop("truncation", False)}
+ mm_kwargs["text_kwargs"] = {"truncation": tok_kwargs.pop("truncation", False)}
+
+ hf_inputs = super()._call_hf_processor(
+ prompt=prompt,
+ mm_data=mm_data,
+ mm_kwargs=mm_kwargs,
+ tok_kwargs=tok_kwargs,
+ )
+
+ if (
+ "audio_feature_lengths" in hf_inputs
+ and "feature_attention_mask" in hf_inputs
+ and (audios := mm_data.get("audio", []))
+ ):
+ audio_num_frames = []
+ for _, audio in enumerate(audios):
+ audio_length = len(audio[0]) if isinstance(audio, tuple) else len(audio)
+ num_frame = (
+ (audio_length // hop_length) if audio_length % hop_length == 0 else (audio_length // hop_length - 1)
+ )
+ if mm_kwargs.get("truncation", False):
+ num_frame = min(num_frame, feature_extractor.n_samples // hop_length)
+ audio_num_frames.append(num_frame)
+ hf_inputs["feature_attention_mask"] = [torch.ones(num_frame) for num_frame in audio_num_frames]
+ hf_inputs["audio_feature_lengths"] = torch.tensor(audio_num_frames)
+ return hf_inputs
+
+ def _maybe_apply_prompt_updates(
+ self,
+ mm_items: MultiModalDataItems,
+ prompt_ids: list[int],
+ mm_kwargs: MultiModalKwargsItems,
+ mm_prompt_updates: MultiModalPromptUpdates,
+ is_update_applied: bool,
+ ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]:
+ """
+ Qwen3-Omni reimplements this function to handle `use_audio_in_video`.
+ """
+ mm_item_counts = mm_items.get_all_counts()
+ self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
+
+ use_audio_in_video = False
+ if "video" in mm_kwargs:
+ for item in mm_kwargs["video"]:
+ if item and item["use_audio_in_video"].data:
+ use_audio_in_video = True
+ else:
+ use_audio_in_video = False
+
+ # normal case with `use_audio_in_video=False`
+ if is_update_applied:
+ mm_placeholders = self._find_mm_placeholders(
+ prompt_ids,
+ mm_prompt_updates,
+ )
+ self._validate_mm_placeholders(
+ mm_placeholders,
+ mm_item_counts,
+ )
+ else:
+ if use_audio_in_video and "audio" in mm_prompt_updates:
+ filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"}
+ prompt_ids, mm_placeholders = self._apply_prompt_updates(
+ prompt_ids,
+ filtered_updates,
+ )
+ # Derive audio placeholders from video placeholders
+ mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates)
+ else:
+ prompt_ids, prompt, mm_placeholders = self._apply_prompt_updates(
+ prompt_ids,
+ mm_prompt_updates,
+ )
+
+ self._validate_mm_placeholders(
+ mm_placeholders,
+ mm_item_counts,
+ )
+
+ return prompt_ids, prompt, mm_placeholders
+
+ def get_updates_use_audio_in_video(
+ self,
+ thinker_config: PretrainedConfig,
+ audio_len: int,
+ video_grid_thw: list[int] | torch.Tensor,
+ video_second_per_grid_t: float,
+ ) -> list[int]:
+ shift = 0
+ audio_token_id = thinker_config.audio_token_id
+ video_token_id = thinker_config.video_token_id
+ audio_start_token_id = thinker_config.audio_start_token_id
+ audio_end_token_id = thinker_config.audio_end_token_id
+ spatial_merge_size = thinker_config.vision_config.spatial_merge_size
+ position_id_per_seconds = thinker_config.position_id_per_seconds
+ audio_token_indices = np.arange(next(iter([audio_len])))
+ curr_video_grid_thw = next(iter([video_grid_thw]))
+ height = curr_video_grid_thw[1] // spatial_merge_size
+ width = curr_video_grid_thw[2] // spatial_merge_size
+ video_token_indices = np.arange(curr_video_grid_thw[0]).reshape(-1, 1, 1)
+ video_token_indices = np.broadcast_to(
+ video_token_indices, (video_token_indices.shape[0], height, width)
+ ).reshape(-1)
+ video_token_indices = (
+ (video_token_indices + shift) * next(iter([video_second_per_grid_t])) * position_id_per_seconds
+ )
+ video_data_index, audio_data_index = 0, 0
+ updates = [audio_start_token_id]
+ while video_data_index < len(video_token_indices) and audio_data_index < len(audio_token_indices):
+ if video_token_indices[video_data_index] <= audio_token_indices[audio_data_index]:
+ updates += [video_token_id]
+ video_data_index += 1
+ else:
+ updates += [audio_token_id]
+ audio_data_index += 1
+ if video_data_index < len(video_token_indices):
+ updates += [video_token_id] * (len(video_token_indices) - video_data_index)
+ if audio_data_index < len(audio_token_indices):
+ updates += [audio_token_id] * (len(audio_token_indices) - audio_data_index)
+ updates += [audio_end_token_id]
+ return updates
+
+ def _get_prompt_updates(
+ self,
+ mm_items: MultiModalDataItems,
+ hf_processor_mm_kwargs: Mapping[str, Any],
+ out_mm_kwargs: MultiModalKwargsItems,
+ ) -> 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]
+
+ out_mm_data = out_mm_kwargs.get_data()
+ audio_feature_lengths = out_mm_data.get("audio_feature_lengths")
+ feature_attention_mask = out_mm_data.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
+ audio_item_idx = 0
+
+ def get_replacement_qwen2_audio(item_idx: int):
+ nonlocal audio_item_idx
+ item_idx += audio_in_video_item_idx
+
+ audio_item_idx += 1
+
+ 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_data[f"{modality}_grid_thw"][item_idx]
+ # 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_item_idx + item_idx]
+ # video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx]
+ audio_num_features = audio_output_lengths[audio_in_video_item_idx + item_idx]
+ video_grid_thw = out_mm_data["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 = 2.0
+
+ placeholder = self.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,
+ )
+ return PromptUpdateDetails.select_token_id(placeholder, embed_token_id=video_token_id)
+
+ 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 _derive_audio_from_video_placeholders(
+ self,
+ placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
+ mm_prompt_updates: Mapping[str, list[PromptUpdateDetails]],
+ ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
+ """
+ Helper to derive audio placeholders from video placeholders when
+ use_audio_in_video=True.
+ """
+ if "video" not in placeholders:
+ return placeholders
+
+ # Validate audio and video counts match
+ num_videos = len(placeholders["video"])
+ num_audios = len(mm_prompt_updates.get("audio", []))
+ if num_audios != num_videos:
+ raise ValueError(
+ f"use_audio_in_video requires equal number of audio and video items, got {num_audios=}, {num_videos=}"
+ )
+
+ tokenizer = self.info.get_tokenizer()
+ processor = self.info.get_hf_processor()
+ audio_token_id = tokenizer.get_vocab()[processor.audio_token]
+
+ result_placeholders = dict(placeholders)
+ audio_placeholders = []
+
+ # Each video is paired with one audio
+ for video_idx, video_placeholder in enumerate(placeholders["video"]):
+ # Create is_embed mask selecting only audio tokens
+ audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id
+
+ audio_placeholder = PlaceholderFeaturesInfo(
+ modality="audio",
+ item_idx=video_idx,
+ start_idx=video_placeholder.start_idx,
+ tokens=video_placeholder.tokens,
+ is_embed=audio_is_embed,
+ )
+ audio_placeholders.append(audio_placeholder)
+
+ result_placeholders["audio"] = audio_placeholders
+ return result_placeholders
+
+ def _get_raw_input_ids(
+ self,
+ token_ids: list[int],
+ use_audio_in_video: bool = False,
+ ) -> list[int]:
+ tokenizer = self.info.get_tokenizer()
+ vision_bos_token = tokenizer.encode(tokenizer.vision_bos_token)[0]
+ vision_eos_token = tokenizer.encode(tokenizer.vision_eos_token)[0]
+ audio_bos_token = tokenizer.encode(tokenizer.audio_bos_token)[0]
+ audio_eos_token = tokenizer.encode(tokenizer.audio_eos_token)[0]
+ audio_token = tokenizer.encode("<|audio_pad|>")[0]
+ image_token = tokenizer.encode("<|image_pad|>")[0]
+ video_token = tokenizer.encode("<|video_pad|>")[0]
+
+ result = token_ids[:]
+ if use_audio_in_video:
+ while True:
+ start = None
+ for i in range(len(result) - 1):
+ if result[i : i + 2] == [vision_bos_token, audio_bos_token]:
+ start = i
+ break
+ if start is not None:
+ end = None
+ for i in range(start + 2, len(result) - 1):
+ if result[i : i + 2] == [audio_eos_token, vision_eos_token]:
+ end = i
+ break
+ if end is not None:
+ result = result[:start] + [vision_bos_token, video_token, vision_eos_token] + result[end + 2 :]
+ else:
+ break
+
+ for mm_token in [audio_token, image_token, video_token]:
+ compressed = []
+ for x in result:
+ if x != mm_token or (not compressed or compressed[-1] != mm_token):
+ compressed.append(x)
+ result = compressed
+
+ return result
+
+
+class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixin):
+ def _process_audio_input(
+ self,
+ audio_input: Qwen2_5OmniAudioFeatureInputs,
+ audio_hashes: list[str] | None = None,
+ cached_audio_features: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ input_features = audio_input["input_features"]
+ audio_feature_lengths = audio_input["audio_feature_lengths"]
+
+ audio_feature_lengths = flatten_bn(audio_feature_lengths, concat=True)
+
+ audio_feat_lengths, audio_output_lengths = _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())
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+ Qwen3OmniMoeThinkerMultiModalProcessor,
+ info=Qwen3OmniMoeThinkerProcessingInfo,
+ dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder,
+)
+class Qwen3OmniMoeThinkerForConditionalGeneration(
+ nn.Module,
+ SupportsMultiModal,
+ SupportsPP,
+ Qwen3OmniMoeConditionalGenerationMixin,
+):
+ merge_by_field_config = True
+
+ 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) -> str | None:
+ if modality.startswith("image"):
+ return "<|vision_start|><|image_pad|><|vision_end|>"
+ if modality.startswith("video"):
+ return "<|vision_start|><|video_pad|><|vision_end|>"
+ if modality.startswith("audio"):
+ return "<|audio_start|><|audio_pad|><|audio_end|>"
+
+ raise ValueError("Only image, video or audio modality is supported")
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ super().__init__()
+ thinker_config: Qwen3OmniMoeThinkerConfig = 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 = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config)
+
+ # attn_backend_override = (
+ # multimodal_config.mm_encoder_attn_backend
+ # if multimodal_config is not None
+ # else None
+ # )
+ self.visual = Qwen3Omni_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"),
+ # attn_backend_override=attn_backend_override,
+ )
+ self.quant_config = quant_config
+
+ lm_vllm_config = vllm_config.with_hf_config(thinker_config.text_config, architectures=["Qwen3MoeForCausalLM"])
+
+ self.language_model = Qwen3MoeLLMForCausalLM(
+ vllm_config=lm_vllm_config,
+ prefix=maybe_prefix(prefix, "language_model"),
+ )
+
+ self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+
+ self.use_deepstack = hasattr(thinker_config.vision_config, "deepstack_visual_indexes")
+ self.deepstack_num_level = (
+ len(thinker_config.vision_config.deepstack_visual_indexes) if self.use_deepstack else 0
+ )
+ # register buffer for deepstack
+ self.deepstack_input_embeds = (
+ [
+ torch.zeros(
+ vllm_config.scheduler_config.max_num_batched_tokens,
+ thinker_config.text_config.hidden_size,
+ )
+ for _ in range(self.deepstack_num_level)
+ ]
+ if self.use_deepstack
+ else None
+ )
+ self.visual_dim = thinker_config.vision_config.out_hidden_size
+ self.multiscale_dim = self.visual_dim * self.deepstack_num_level
+
+ def _get_deepstack_input_embeds(self, num_tokens: int) -> IntermediateTensors:
+ # get deepstack_input_embeds from buffer, and clear the buffer
+ return IntermediateTensors(
+ {
+ f"deepstack_input_embeds_{idx}": self.deepstack_input_embeds[idx][:num_tokens]
+ for idx in range(self.deepstack_num_level)
+ }
+ )
+
+ def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> None:
+ # set deepstack_input_embeds to buffer
+ num_tokens = deepstack_input_embeds.size(1)
+ if num_tokens > self.deepstack_input_embeds[0].size(0):
+ self.deepstack_input_embeds = [
+ torch.zeros(
+ num_tokens,
+ self.config.text_config.hidden_size,
+ device=self.deepstack_input_embeds[0].device,
+ dtype=self.deepstack_input_embeds[0].dtype,
+ )
+ for _ in range(self.deepstack_num_level)
+ ]
+ for idx in range(self.deepstack_num_level):
+ self.deepstack_input_embeds[idx][:num_tokens].copy_(deepstack_input_embeds[idx])
+
+ def _clear_deepstack_input_embeds(self, num_tokens: int) -> None:
+ # clear deepstack_input_embeds in buffer
+ if num_tokens > 0:
+ for idx in range(self.deepstack_num_level):
+ self.deepstack_input_embeds[idx][:num_tokens].zero_()
+
+ 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 | None:
+ 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":
+ image_embeddings = self._process_image_input(multimodal_input)
+ multimodal_embeddings += tuple(image_embeddings)
+ if modality == "video":
+ video_embeddings = self._process_video_input(multimodal_input)
+ multimodal_embeddings += tuple(video_embeddings)
+ if modality == "audio":
+ audio_embeddings = self._process_audio_input(multimodal_input)
+ multimodal_embeddings += tuple(audio_embeddings)
+ return multimodal_embeddings
+
+ def _get_text_embeddings(
+ self,
+ input_ids: Tensor,
+ embed_input_ids: Callable[[Tensor], Tensor],
+ *,
+ is_multimodal: Tensor | None,
+ handle_oov_mm_token: bool,
+ ) -> Tensor:
+ if handle_oov_mm_token and is_multimodal is not None:
+ is_text = ~is_multimodal
+ text_embeds = embed_input_ids(input_ids[is_text])
+
+ return torch.empty(
+ (input_ids.shape[0], text_embeds.shape[1]),
+ dtype=text_embeds.dtype,
+ device=text_embeds.device,
+ ).masked_scatter_(is_text.unsqueeze_(-1), text_embeds)
+
+ return embed_input_ids(input_ids)
+
+ def get_input_embeddings(
+ self,
+ input_ids: torch.Tensor,
+ multimodal_embeddings: MultiModalEmbeddings | None = None,
+ *,
+ is_multimodal: torch.Tensor | None = None,
+ handle_oov_mm_token: bool = False,
+ ) -> torch.Tensor:
+ # Compute is_multimodal if not provided and we have multimodal embeddings
+ if is_multimodal is None and multimodal_embeddings is not None and len(multimodal_embeddings) > 0:
+ # Create a boolean tensor marking which positions are multimodal tokens
+ is_multimodal = (
+ (input_ids == self.config.image_token_id)
+ | (input_ids == self.config.video_token_id)
+ | (input_ids == self.config.audio_token_id)
+ )
+
+ inputs_embeds = self._get_text_embeddings(
+ input_ids,
+ self.language_model.get_input_embeddings,
+ is_multimodal=is_multimodal,
+ handle_oov_mm_token=handle_oov_mm_token,
+ )
+
+ if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+ return inputs_embeds
+
+ deepstack_input_embeds = None
+ # TODO (ywang96): support overlapping modalitiy embeddings so that
+ # `use_audio_in_video` will work on V1.
+ # split the feat dim to obtain multi-scale visual feature
+ has_vision_embeddings = [
+ embeddings.shape[-1] != self.config.text_config.hidden_size for embeddings in multimodal_embeddings
+ ]
+ if self.visual.deepstack_visual_indexes is not None and any(has_vision_embeddings):
+ multiscale_len = len(self.visual.deepstack_visual_indexes)
+ multimodal_embeddings_multiscale = []
+ is_vision = torch.zeros_like(is_multimodal)
+ mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0]
+ mm_position_idx = 0
+ for index, embeddings in enumerate(multimodal_embeddings):
+ num_tokens = embeddings.shape[0]
+ current_positions = mm_positions[mm_position_idx : mm_position_idx + num_tokens]
+
+ # Vision embeddings
+ if embeddings.shape[-1] != self.config.text_config.hidden_size:
+ visual_dim = embeddings.shape[-1] // (multiscale_len + 1)
+ multi_dim = visual_dim * multiscale_len
+ embeddings_main, embeddings_multiscale = torch.split(embeddings, [visual_dim, multi_dim], dim=-1)
+ multimodal_embeddings[index] = embeddings_main
+ multimodal_embeddings_multiscale.append(embeddings_multiscale)
+ is_vision[current_positions] = True
+
+ # Audio embeddings
+ else:
+ is_vision[current_positions] = False
+
+ mm_position_idx += num_tokens
+
+ deepstack_input_embeds = inputs_embeds.new_zeros(
+ inputs_embeds.size(0), multiscale_len * inputs_embeds.size(1)
+ )
+ deepstack_input_embeds = _merge_multimodal_embeddings(
+ inputs_embeds=deepstack_input_embeds,
+ multimodal_embeddings=multimodal_embeddings_multiscale,
+ is_multimodal=is_vision,
+ )
+ deepstack_input_embeds = (
+ deepstack_input_embeds.view(inputs_embeds.shape[0], multiscale_len, visual_dim)
+ .permute(1, 0, 2)
+ .contiguous()
+ )
+ self._set_deepstack_input_embeds(deepstack_input_embeds)
+
+ inputs_embeds = _merge_multimodal_embeddings(
+ inputs_embeds=inputs_embeds,
+ multimodal_embeddings=multimodal_embeddings,
+ is_multimodal=is_multimodal,
+ )
+
+ return inputs_embeds
+
+ def forward(
+ self,
+ input_ids: torch.Tensor,
+ positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ *,
+ capture_layer_indices: Sequence[int] | None = None,
+ return_hidden_states: bool = False,
+ **kwargs: object,
+ ) -> torch.Tensor | IntermediateTensors:
+ if intermediate_tensors is not None:
+ inputs_embeds = None
+
+ if self.use_deepstack and inputs_embeds is not None and get_pp_group().is_first_rank:
+ deepstack_input_embeds = self._get_deepstack_input_embeds(inputs_embeds.size(0))
+ else:
+ deepstack_input_embeds = None
+
+ hidden_states, captured_hidden_states = self.language_model.model(
+ input_ids,
+ positions,
+ intermediate_tensors,
+ inputs_embeds=inputs_embeds,
+ capture_layer_indices=capture_layer_indices,
+ return_hidden_states=return_hidden_states,
+ # args for deepstack
+ deepstack_input_embeds=deepstack_input_embeds,
+ )
+
+ if inputs_embeds is not None and get_pp_group().is_first_rank:
+ self._clear_deepstack_input_embeds(inputs_embeds.size(0))
+
+ return hidden_states.unsqueeze(0), captured_hidden_states
+
+ def compute_logits(
+ self,
+ hidden_states: torch.Tensor,
+ sampling_metadata: SamplingMetadata = None,
+ ) -> torch.Tensor | None:
+ 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.", "code2wav."],
+ )
+ 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
+
+ def get_mrope_input_positions(
+ self,
+ input_tokens: list[int],
+ hf_config: PretrainedConfig,
+ image_grid_thw: list[list[int]] | torch.Tensor | None,
+ video_grid_thw: list[list[int]] | torch.Tensor | None,
+ second_per_grid_ts: list[float] | None = None,
+ context_len: int = 0,
+ seq_len: int | None = None,
+ audio_feature_lengths: torch.Tensor | None = None,
+ use_audio_in_video: bool = False,
+ ) -> tuple[torch.Tensor, int]:
+ config = hf_config.thinker_config
+ 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)
+ input_ids = torch.tensor(input_tokens)
+ if input_ids is None or input_ids.ndim != 1:
+ raise ValueError("_omni3_get_input_positions_tensor expects 1D input_ids")
+
+ seq_len = input_ids.shape[0]
+ if audio_feature_lengths is not None and not isinstance(audio_feature_lengths, torch.Tensor):
+ audio_feature_lengths = torch.as_tensor(audio_feature_lengths, dtype=torch.long)
+ if not len(second_per_grid_ts) and len(video_grid_thw):
+ second_per_grid_ts = 2.0
+ second_per_grids = torch.ones(len(video_grid_thw), dtype=torch.float32) * second_per_grid_ts
+ else:
+ second_per_grids = torch.tensor(second_per_grid_ts, dtype=torch.float32)
+
+ spatial_merge_size = config.vision_config.spatial_merge_size
+ image_token_id = config.image_token_id
+ video_token_id = config.video_token_id
+ audio_token_id = config.audio_token_id
+ vision_start_token_id = config.vision_start_token_id
+ audio_start_token_id = config.audio_start_token_id
+ position_id_per_seconds = config.position_id_per_seconds
+
+ vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+ if vision_start_indices.numel() > 0:
+ vision_tokens = input_ids[vision_start_indices + 1]
+ else:
+ vision_tokens = input_ids.new_empty((0,), dtype=input_ids.dtype)
+ audio_nums = torch.sum(input_ids == audio_start_token_id)
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (
+ (vision_tokens == audio_start_token_id).sum()
+ if use_audio_in_video
+ else (vision_tokens == video_token_id).sum()
+ )
+
+ llm_pos_ids_list: list[torch.Tensor] = []
+ st = 0
+ image_idx = 0
+ video_idx = 0
+ audio_idx = 0
+ remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums # noqa: E501
+ multimodal_nums = image_nums + audio_nums if use_audio_in_video else image_nums + video_nums + audio_nums # noqa: E501
+
+ for _ in range(multimodal_nums):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ if (image_token_id in input_tokens or video_token_id in input_tokens) and (
+ remain_videos > 0 or remain_images > 0
+ ):
+ ed_vision_start = input_tokens.index(vision_start_token_id, st)
+ else:
+ ed_vision_start = len(input_tokens) + 1
+ if audio_token_id in input_tokens and remain_audios > 0:
+ ed_audio_start = input_tokens.index(audio_start_token_id, st)
+ else:
+ ed_audio_start = len(input_tokens) + 1
+ min_ed = min(ed_vision_start, ed_audio_start)
+
+ if min_ed == ed_audio_start:
+ text_len = min_ed - st
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ _, audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx])
+ llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx
+ llm_pos_ids_list.append(llm_pos_ids)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+ st += text_len + bos_len + audio_len + eos_len
+ audio_idx += 1
+ remain_audios -= 1
+ elif min_ed == ed_vision_start and input_ids[ed_vision_start + 1] == image_token_id:
+ text_len = min_ed - st
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ 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) * position_id_per_seconds
+ llm_pos_ids = get_llm_pos_ids_for_vision(
+ st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
+ llm_pos_ids_list.append(llm_pos_ids)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+ st += text_len + bos_len + image_len + eos_len
+ image_idx += 1
+ remain_images -= 1
+ elif (
+ min_ed == ed_vision_start
+ and input_ids[ed_vision_start + 1] == video_token_id
+ and not use_audio_in_video
+ ):
+ text_len = min_ed - st
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ 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) * float(second_per_grids[video_idx].item()) * position_id_per_seconds
+ llm_pos_ids = get_llm_pos_ids_for_vision(
+ st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+ llm_pos_ids_list.append(llm_pos_ids)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+ st += text_len + bos_len + video_len + eos_len
+ video_idx += 1
+ remain_videos -= 1
+ elif min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start and use_audio_in_video:
+ text_len = min_ed - st
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ bos_len = 1
+ bos_block = torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx
+ llm_pos_ids_list.append(bos_block)
+ llm_pos_ids_list.append(bos_block)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ _, audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx])
+ audio_llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx
+ 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) * float(second_per_grids[video_idx].item()) * position_id_per_seconds
+ video_llm_pos_ids = get_llm_pos_ids_for_vision(
+ st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ video_data_index, audio_data_index = 0, 0
+ while video_data_index < video_llm_pos_ids.shape[-1] and audio_data_index < audio_llm_pos_ids.shape[-1]:
+ if video_llm_pos_ids[0][video_data_index] <= audio_llm_pos_ids[0][audio_data_index]:
+ llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_data_index + 1])
+ video_data_index += 1
+ else:
+ llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_data_index + 1])
+ audio_data_index += 1
+ if video_data_index < video_llm_pos_ids.shape[-1]:
+ llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_llm_pos_ids.shape[-1]])
+ if audio_data_index < audio_llm_pos_ids.shape[-1]:
+ llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_llm_pos_ids.shape[-1]])
+ video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ eos_len = 1
+ eos_block = torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx
+ llm_pos_ids_list.append(eos_block)
+ llm_pos_ids_list.append(eos_block)
+ st += text_len + bos_len * 2 + audio_len + video_len + eos_len * 2 # noqa: E501
+ audio_idx += 1
+ video_idx += 1
+ remain_videos -= 1
+ remain_audios -= 1
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+ if llm_positions.shape[1] != seq_len:
+ raise RuntimeError("Position ids length mismatch with input ids length")
+
+ mrope_position_delta = llm_positions.max() + 1 - seq_len
+ return llm_positions, mrope_position_delta
diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py
index 8e6977950..fd64b8b2e 100644
--- a/vllm_omni/model_executor/models/registry.py
+++ b/vllm_omni/model_executor/models/registry.py
@@ -27,6 +27,27 @@ _OMNI_MODELS = {
"Qwen2_5OmniToken2WavModel",
),
"Qwen2ForCausalLM_old": ("qwen2_5_omni", "qwen2_old", "Qwen2ForCausalLM"), # need to discuss
+ # Qwen3 Omni MoE models
+ "Qwen3OmniMoeForConditionalGeneration": (
+ "qwen3_omni",
+ "qwen3_omni",
+ "Qwen3OmniMoeForConditionalGeneration",
+ ),
+ "Qwen3OmniMoeThinkerForConditionalGeneration": (
+ "qwen3_omni",
+ "qwen3_omni_moe_thinker",
+ "Qwen3OmniMoeThinkerForConditionalGeneration",
+ ),
+ "Qwen3OmniMoeTalkerForConditionalGeneration": (
+ "qwen3_omni",
+ "qwen3_omni_moe_talker",
+ "Qwen3OmniMoeTalkerForConditionalGeneration",
+ ),
+ "Qwen3OmniMoeCode2Wav": (
+ "qwen3_omni",
+ "qwen3_omni_code2wav",
+ "Qwen3OmniMoeCode2Wav",
+ ),
}
_VLLM_OMNI_MODELS = {
diff --git a/vllm_omni/model_executor/models/utils.py b/vllm_omni/model_executor/models/utils.py
index f3854ad4e..7f73067b3 100644
--- a/vllm_omni/model_executor/models/utils.py
+++ b/vllm_omni/model_executor/models/utils.py
@@ -28,3 +28,12 @@ def split_list_into_ranges(lst: torch.Tensor, interval: int) -> list[list[int]]:
ranges[index].append(num)
return ranges
+
+
+def safe_tensor_reshape(tensor: torch.Tensor, shape: tuple) -> torch.Tensor:
+ """
+ Reshape a tensor safely.
+ """
+ if tensor is None:
+ return None
+ return tensor.reshape(shape)
diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
index 54d2b2359..c378a5b48 100644
--- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
@@ -11,9 +11,9 @@ stage_args:
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
gpu_memory_utilization: 0.8
- enforce_eager: true # need to discuss
+ enforce_eager: true # Now we only support eager mode
trust_remote_code: true
- engine_output_type: latent # change the param name,such as pooling_output
+ engine_output_type: latent
enable_prefix_caching: false
is_comprehension: true
final_output: true
@@ -30,7 +30,7 @@ stage_args:
runtime:
process: true
devices: "1"
- max_batch_size: 3
+ max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
new file mode 100644
index 000000000..8584b4c6d
--- /dev/null
+++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml
@@ -0,0 +1,94 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
+stage_args:
+ - stage_id: 0
+ runtime:
+ devices: "0"
+ max_batch_size: 1
+ engine_args:
+ model_stage: thinker
+ model_arch: Qwen3OmniMoeForConditionalGeneration
+ worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+ scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+ gpu_memory_utilization: 0.8
+ enforce_eager: true
+ trust_remote_code: true
+ engine_output_type: latent # Output hidden states for talker
+ distributed_executor_backend: "mp"
+ enable_prefix_caching: false
+ hf_config_name: thinker_config
+ # tensor_parallel_size: 2 # Commented out - using single GPU for thinker
+ final_output: true
+ final_output_type: text
+ is_comprehension: true
+ default_sampling_params:
+ temperature: 0.4
+ top_p: 0.9
+ top_k: 1
+ max_tokens: 2048
+ seed: 42
+ detokenize: True
+ repetition_penalty: 1.05
+
+ - stage_id: 1
+ runtime:
+ devices: "1"
+ max_batch_size: 1
+ engine_args:
+ model_stage: talker
+ model_arch: Qwen3OmniMoeForConditionalGeneration
+ worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+ scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+ gpu_memory_utilization: 0.5
+ enforce_eager: true
+ trust_remote_code: true
+ engine_output_type: latent # Output codec codes for code2wav
+ # tensor_parallel_size: 2
+ enable_prefix_caching: false
+ distributed_executor_backend: "mp"
+ hf_config_name: talker_config
+ engine_input_source: [0]
+ custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+ # final_output: true
+ # final_output_type: text
+ default_sampling_params:
+ temperature: 0.9
+ top_k: 50
+ max_tokens: 4096
+ seed: 42
+ detokenize: False
+ repetition_penalty: 1.05
+ stop_token_ids: [2150]
+
+ - stage_id: 2
+ runtime:
+ devices: "1"
+ max_batch_size: 1
+ engine_args:
+ model_stage: code2wav
+ model_arch: Qwen3OmniMoeForConditionalGeneration
+ worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+ scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler
+ enforce_eager: true
+ trust_remote_code: true
+ enable_prefix_caching: false
+ engine_output_type: audio # Final output: audio waveform
+ gpu_memory_utilization: 0.1
+ distributed_executor_backend: "mp"
+ max_num_batched_tokens: 1000000
+ hf_config_name: thinker_config
+ engine_input_source: [1]
+ custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+ final_output: true
+ final_output_type: audio
+ default_sampling_params:
+ temperature: 0.0
+ top_p: 1.0
+ top_k: -1
+ max_tokens: 65536
+ seed: 42
+ detokenize: True
+ repetition_penalty: 1.1
diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
new file mode 100644
index 000000000..e2d07aa1b
--- /dev/null
+++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright 2025 The Qwen team.
+"""Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition."""
+
+from typing import Union
+
+import torch
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+
+def _compute_talker_prompt_ids_length(info):
+ im_start_token_id = 151644
+ system_token_id = 8948
+ user_token_id = 872
+ assistant_token_id = 77091
+
+ thinker_sequences = torch.tensor(info["thinker_sequences"]).unsqueeze(0).cuda().long() # [1, T]
+ input_ids = torch.tensor(info["thinker_input_ids"]).unsqueeze(0).cuda().long() # [1, T]
+
+ im_start_indexes = torch.cat(
+ [
+ torch.nonzero(input_ids[0] == im_start_token_id).squeeze(1),
+ torch.tensor([thinker_sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
+ ],
+ dim=0,
+ )
+
+ sum_user_len = 0
+ assistant_len = 0
+ for i in range(len(im_start_indexes) - 1):
+ s = int(im_start_indexes[i].item())
+ e = int(im_start_indexes[i + 1].item())
+ role = int(input_ids[0, s + 1].item())
+ if role == system_token_id:
+ continue
+ elif role == user_token_id:
+ sum_user_len += e - s
+ elif role == assistant_token_id and i == len(im_start_indexes) - 2:
+ assistant_len += 9 # 3 + 4 + 1 + 1
+ else:
+ pass
+
+ return sum_user_len + assistant_len
+
+
+def thinker2talker(
+ stage_list,
+ engine_input_source,
+ prompt: Union[OmniTokensPrompt, TextPrompt] = None,
+ requires_multimodal_data: bool = False,
+):
+ """
+ Process thinker outputs to create talker inputs.
+
+ Workflow:
+ 1. Extract thinker's text generation outputs (token IDs + hidden states)
+ 2. Split hidden states into: prompt embeddings + generated embeddings
+ 3. Package for talker with additional information
+
+ Args:
+ stage_list: List of stage objects
+ engine_input_source: Source stage IDs (typically [0] for thinker)
+ prompt: Original prompt data
+
+ Returns:
+ List of OmniTokensPrompt for talker stage
+ """
+ if not engine_input_source:
+ raise ValueError("engine_input_source cannot be empty")
+
+ source_stage_id = engine_input_source[0]
+ if source_stage_id >= len(stage_list):
+ raise IndexError(f"Invalid stage_id: {source_stage_id}")
+
+ if stage_list[source_stage_id].engine_outputs is None:
+ raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
+
+ thinker_outputs = stage_list[source_stage_id].engine_outputs
+ talker_inputs = []
+
+ # Process each thinker output
+ for i, thinker_output in enumerate(thinker_outputs):
+ output = thinker_output.outputs[0]
+ thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
+
+ thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
+ info = {
+ "thinker_embeddings": thinker_embeddings,
+ "thinker_hidden_states": thinker_hidden_states,
+ "thinker_sequences": thinker_output.prompt_token_ids
+ + output.token_ids, # the thinker_sequences is the whole ids
+ "thinker_input_ids": thinker_output.prompt_token_ids,
+ # Provide thinker-side TTS token embeddings for talker projection
+ "tts_bos_embed": output.multimodal_output.get("tts_bos_embed").float().clone().detach().cuda(),
+ "tts_eos_embed": output.multimodal_output.get("tts_eos_embed").float().clone().detach().cuda(),
+ "tts_pad_embed": output.multimodal_output.get("tts_pad_embed").float().clone().detach().cuda(),
+ }
+ talker_inputs.append(
+ OmniTokensPrompt(
+ prompt_token_ids=[0] * _compute_talker_prompt_ids_length(info),
+ additional_information=info,
+ multi_modal_data=None,
+ mm_processor_kwargs=None,
+ )
+ )
+
+ return talker_inputs
+
+
+def talker2code2wav(
+ stage_list,
+ engine_input_source,
+ prompt: Union[OmniTokensPrompt, TextPrompt] = None,
+ requires_multimodal_data: bool = False,
+):
+ """
+ Process talker outputs to create code2wav inputs.
+
+ Workflow:
+ 1. Extract talker's codec code outputs (8-layer RVQ codes)
+ 2. Flatten codes for code2wav input
+ 3. Package for code2wav stage
+
+ Args:
+ stage_list: List of stage objects
+ engine_input_source: Source stage IDs (typically [1] for talker)
+ prompt: Original prompt data
+
+ Returns:
+ List of OmniTokensPrompt for code2wav stage
+ """
+ if not engine_input_source:
+ raise ValueError("engine_input_source cannot be empty")
+
+ source_stage_id = engine_input_source[0]
+ if source_stage_id >= len(stage_list):
+ raise IndexError(f"Invalid stage_id: {source_stage_id}")
+
+ if stage_list[source_stage_id].engine_outputs is None:
+ raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
+
+ talker_outputs = stage_list[source_stage_id].engine_outputs
+ code2wav_inputs = []
+
+ # Process each talker output
+ for i, talker_output in enumerate(talker_outputs):
+ output = talker_output.outputs[0]
+
+ # Extract codec codes from talker output
+ # Expected shape: [8, seq_len] (8-layer RVQ codes)
+ codec_codes = (
+ output.multimodal_output["code_predictor_codes"]
+ .to(torch.long)
+ .transpose(0, 1)
+ .cpu()
+ .to(torch.long)
+ .reshape(-1)
+ .tolist()
+ ) # 16, seq_len
+ code2wav_inputs.append(
+ OmniTokensPrompt(
+ prompt_token_ids=codec_codes,
+ multi_modal_data=None,
+ mm_processor_kwargs=None,
+ )
+ )
+
+ return code2wav_inputs
diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py
index 454e7b423..3e1e2e036 100644
--- a/vllm_omni/outputs.py
+++ b/vllm_omni/outputs.py
@@ -7,11 +7,34 @@ from vllm.v1.outputs import ModelRunnerOutput
class OmniModelRunnerOutput(ModelRunnerOutput):
+ """Model runner output for omni models.
+
+ Extends the base ModelRunnerOutput with support for multimodal outputs
+ that may be produced by non-autoregressive stages.
+
+ Attributes:
+ multimodal_outputs: Optional dictionary mapping modality names to
+ output tensors (e.g., {"image": tensor, "audio": tensor})
+ """
+
multimodal_outputs: Optional[dict[str, torch.Tensor]] = None
@dataclass
class OmniRequestOutput(RequestOutput):
+ """Request output for omni pipeline stages.
+
+ Wraps a standard RequestOutput with stage-specific metadata,
+ indicating which stage produced the output and what type of
+ final output it represents.
+
+ Attributes:
+ stage_id: Identifier of the stage that produced this output
+ final_output_type: Type of final output (e.g., "text", "image",
+ "audio", "latents")
+ request_output: The underlying RequestOutput from the stage
+ """
+
stage_id: int
final_output_type: str
request_output: RequestOutput
diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py
index ea0828cff..a0028d572 100644
--- a/vllm_omni/patch.py
+++ b/vllm_omni/patch.py
@@ -1,16 +1,24 @@
import sys
from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
-from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding as _OriginalMRotaryEmbedding
+from vllm.model_executor.layers.rotary_embedding import (
+ MRotaryEmbedding as _OriginalMRotaryEmbedding,
+)
+from vllm.v1.engine import EngineCoreOutput as _OriginalEngineCoreOutput
+from vllm.v1.engine import EngineCoreOutputs as _OriginalEngineCoreOutputs
from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest
from vllm.v1.request import Request as _OriginalRequest
-from vllm_omni.engine import OmniEngineCoreRequest
+from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEngineCoreRequest
from vllm_omni.inputs.data import OmniTokensPrompt
from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
from vllm_omni.request import OmniRequest
for module_name, module in sys.modules.items():
+ if hasattr(module, "EngineCoreOutput") and module.EngineCoreOutput == _OriginalEngineCoreOutput:
+ module.EngineCoreOutput = OmniEngineCoreOutput
+ if hasattr(module, "EngineCoreOutputs") and module.EngineCoreOutputs == _OriginalEngineCoreOutputs:
+ module.EngineCoreOutputs = OmniEngineCoreOutputs
if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt:
module.TokensPrompt = OmniTokensPrompt
if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding:
diff --git a/vllm_omni/request.py b/vllm_omni/request.py
index bc71bf968..31f90f94b 100644
--- a/vllm_omni/request.py
+++ b/vllm_omni/request.py
@@ -10,6 +10,19 @@ from vllm_omni.engine import AdditionalInformationPayload, OmniEngineCoreRequest
class OmniRequest(Request):
+ """Request class for omni models, extending the base Request.
+
+ This class extends the base vLLM Request with support for prompt
+ embeddings and additional information payloads, enabling direct
+ transfer of pre-computed embeddings between stages.
+
+ Args:
+ prompt_embeds: Optional serialized prompt embeddings payload.
+ Used for direct transfer of embeddings between stages.
+ additional_information: Optional additional information payload
+ containing tensors or lists to be passed along with the request.
+ """
+
def __init__(
self,
prompt_embeds: Optional[PromptEmbedsPayload] = None,
@@ -29,6 +42,16 @@ class OmniRequest(Request):
request: OmniEngineCoreRequest,
block_hasher: Optional[Callable[["Request"], list["BlockHash"]]],
) -> "Request":
+ """Create an OmniRequest from an OmniEngineCoreRequest.
+
+ Args:
+ request: The OmniEngineCoreRequest to convert
+ block_hasher: Optional function to compute block hashes for
+ prefix caching
+
+ Returns:
+ OmniRequest instance created from the engine core request
+ """
return cls(
request_id=request.request_id,
client_index=request.client_index,
diff --git a/vllm_omni/version.py b/vllm_omni/version.py
index 92007aa40..d66839fdf 100644
--- a/vllm_omni/version.py
+++ b/vllm_omni/version.py
@@ -1,12 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-try:
- from ._version import __version__, __version_tuple__
-except Exception as e:
- import warnings
-
- warnings.warn(f"Failed to read commit hash:\n{e}", RuntimeWarning, stacklevel=2)
-
- __version__ = "dev"
- __version_tuple__ = (0, 0, __version__)
+__version__ = "0.11.0rc1"
+__version_tuple__ = (0, 11, "0rc1")
+# TODO: add auto version generation
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index a0f6780ca..8a0d28785 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -1,4 +1,4 @@
-"""AR GPU Model Runner for vLLM-omni.
+"""AR GPU Model Runner for vLLM-Omni.
Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
and also outputs sampled tokens.
@@ -6,9 +6,6 @@ and also outputs sampled tokens.
from __future__ import annotations
-from typing import Any
-
-import numpy as np
import torch
from vllm.forward_context import BatchDescriptor
from vllm.logger import init_logger
@@ -25,10 +22,8 @@ from vllm.v1.worker.gpu_model_runner import (
has_kv_transfer_group,
set_forward_context,
)
-from vllm.v1.worker.ubatch_utils import UBatchSlices
from vllm.v1.worker.utils import is_residual_scattered_for_sp
-from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
from vllm_omni.outputs import OmniModelRunnerOutput
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
@@ -48,168 +43,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
-
- def _preprocess(
- self,
- scheduler_output: SchedulerOutput,
- num_scheduled_tokens_np: np.ndarray,
- intermediate_tensors: IntermediateTensors | None = None,
- ubatch_slices: UBatchSlices | None = None,
- num_tokens_after_padding: torch.Tensor | None = None,
- ) -> tuple[
- int,
- int,
- torch.Tensor | None,
- torch.Tensor | None,
- torch.Tensor | None,
- torch.Tensor,
- IntermediateTensors | None,
- dict[str, Any],
- dict[str, dict] | None,
- ]:
- num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
- if ubatch_slices:
- assert num_tokens_after_padding is not None
- num_input_tokens = int(num_tokens_after_padding[0].item() * 2)
- self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
- elif ubatch_slices is None:
- num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)
- num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
- num_input_tokens += num_pad
-
- # _prepare_inputs may reorder the batch, so we must gather multi
- # modal outputs after that to ensure the correct order
- if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
- # Run the multimodal encoder if any.
- self._execute_mm_encoder(scheduler_output)
- mm_embeds = self._gather_mm_embeddings(scheduler_output)
-
- # NOTE(woosuk): To unify token ids and soft tokens (vision
- # embeddings), we always use embeddings (rather than token ids)
- # as input to the multimodal model, even when the input is text.
- inputs_embeds_scheduled = self.model.get_input_embeddings(
- input_ids=self.input_ids.gpu[:num_scheduled_tokens],
- multimodal_embeddings=mm_embeds or None,
- )
-
- # TODO(woosuk): Avoid the copy. Optimize.
- self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
-
- # Reset per-step additional information collector (deprecated concat path)
- if hasattr(self, "_forward_additional_information"):
- self._forward_additional_information = None
- # New: per-request additional information for this step
-
- # Overlay custom prompt_embeds per request for the prompt portion;
- # collect additional_information (tensor/list) for prefill portion only
- for req_index, req_id in enumerate(self.input_batch.req_ids):
- req_state = self.requests[req_id]
- pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
- num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index])
- prompt_len = len(req_state.prompt_token_ids)
- prompt_remaining = max(0, prompt_len - num_computed_tokens)
- sched_tokens = int(num_scheduled_tokens_np[req_index])
- overlay_len = min(sched_tokens, prompt_remaining)
- if overlay_len <= 0:
- continue
- if pe_cpu is not None:
- src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to(
- dtype=self.dtype, device=self.device, non_blocking=True
- )
- start_offset = int(self.query_start_loc.cpu[req_index])
- self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
-
- input_ids = self.input_ids.gpu[:num_input_tokens]
- inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
- model_kwargs = {
- **self._init_model_kwargs(num_scheduled_tokens),
- **self._extract_mm_kwargs(scheduler_output),
- }
- elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
- # Get the input embeddings for the tokens that are not input embeds,
- # then put them into the appropriate positions.
- # TODO(qthequartermasterman): Since even when prompt embeds are
- # enabled, (a) not all requests will use prompt embeds, and (b)
- # after the initial prompt is processed, the rest of the generated
- # tokens will be token ids, it is not desirable to have the
- # embedding layer outside of the CUDA graph all the time. The v0
- # engine avoids this by "double compiling" the CUDA graph, once
- # with input_ids and again with inputs_embeds, for all num_tokens.
- # If a batch only has token ids, then including the embedding layer
- # in the CUDA graph will be more performant (like in the else case
- # below).
- token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1)
- # Some tokens ids may need to become embeds
- if token_ids_idx.numel() > 0:
- token_ids = self.input_ids.gpu[token_ids_idx]
- tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
- self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
-
- inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
- model_kwargs = self._init_model_kwargs(num_input_tokens)
- input_ids = None
- else:
- # For text-only models, we use token ids as input.
- # While it is possible to use embeddings as input just like the
- # multimodal models, it is not desirable for performance since
- # then the embedding layer is not included in the CUDA graph.
- input_ids = self.input_ids.gpu[:num_input_tokens]
- inputs_embeds = None
- model_kwargs = self._init_model_kwargs(num_input_tokens)
- if self.uses_mrope:
- positions = self.mrope_positions.gpu[:, :num_input_tokens]
- else:
- positions = self.positions.gpu[:num_input_tokens]
-
- if get_pp_group().is_first_rank:
- intermediate_tensors = None
- else:
- intermediate_tensors = self.sync_and_slice_intermediate_tensors(
- num_input_tokens, intermediate_tensors, True
- )
-
- if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
- encoder_inputs = self._extract_encoder_inputs(scheduler_output)
- model_kwargs.update(encoder_inputs)
-
- if self.model.has_preprocess:
- # Overlay custom prompt_embeds per request for the prompt portion;
- # collect additional_information (tensor/list) for prefill portion only
- for req_index, req_id in enumerate(self.input_batch.req_ids):
- req_state = self.requests.get(req_id)
- req_infos = (getattr(req_state, "additional_information_cpu", None)
- if req_state is not None else None)
-
- start_offset = int(self.query_start_loc.cpu[req_index])
- sched_tokens = int(num_scheduled_tokens_np[req_index])
- s, e = start_offset, start_offset + sched_tokens
- span_len = int(e) - int(s)
-
- # call the custom process function
- req_input_ids, req_embeds, update_dict = self.model.preprocess(
- input_ids=input_ids[s:e],
- inputs_embeds=inputs_embeds[s:e],
- **req_infos
- )
- # TODO(Peiqi): the merge stage could move out from the critical path
- self._merge_additional_information_update(req_id, update_dict)
-
- # update the inputs_embeds and input_ids
- 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
-
- return (
- num_scheduled_tokens,
- num_input_tokens,
- num_tokens_after_padding,
- input_ids,
- inputs_embeds,
- positions,
- intermediate_tensors,
- model_kwargs,
- )
+ # each model stage has their own hidden size
+ self.hidden_size = self.model_config.hf_text_config.hidden_size
+ self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
@torch.inference_mode()
def execute_model(
@@ -222,58 +58,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
# Update persistent batch states.
self._update_states(scheduler_output)
- # Decode per-request prompt_embeds / additional_hidden_states payloads
+ # Decode per-request prompt_embeds / additional_information payloads
# (if present) into CPU tensors
- try:
- new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
- if new_reqs:
- for nr in new_reqs:
- req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
- if req_id is None:
- continue
- # prompt_embeds
- payload_pe = getattr(nr, "prompt_embeds", None)
- if payload_pe is not None:
- if isinstance(payload_pe, torch.Tensor):
- pe_cpu = payload_pe.detach().to("cpu").contiguous()
- elif isinstance(payload_pe, PromptEmbedsPayload):
- dt = np.dtype(getattr(payload_pe, "dtype", "float32"))
- arr = np.frombuffer(payload_pe.data, dtype=dt)
- arr = arr.reshape(payload_pe.shape)
- pe_cpu = torch.from_numpy(arr)
- else:
- pe_cpu = None
- if pe_cpu is not None and req_id in self.requests:
- setattr(
- self.requests[req_id],
- "prompt_embeds_cpu",
- pe_cpu,
- )
- # additional_information
- payload_info = getattr(nr, "additional_information", None)
- if payload_info is not None:
- info_dict = {}
- if isinstance(payload_info, dict):
- # Already decoded
- info_dict = payload_info
- elif isinstance(payload_info, AdditionalInformationPayload):
- for k, entry in payload_info.entries.items():
- if entry.tensor_data is not None:
- dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
- arr = np.frombuffer(entry.tensor_data, dtype=dt)
- arr = arr.reshape(entry.tensor_shape)
- info_dict[k] = torch.from_numpy(arr)
- else:
- info_dict[k] = entry.list_data
- if info_dict and req_id in self.requests:
- setattr(
- self.requests[req_id],
- "additional_information_cpu",
- info_dict,
- )
- except Exception as e:
- logger.error(f"Error decoding prompt_embeds / additional_information: {e}")
- pass
+ self._decode_and_store_request_payloads(scheduler_output)
if not scheduler_output.total_num_scheduled_tokens:
if not has_kv_transfer_group():
@@ -308,6 +95,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
positions,
intermediate_tensors,
model_kwargs,
+ per_req_additional_information,
) = self._preprocess(
scheduler_output,
num_scheduled_tokens_np,
@@ -342,9 +130,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
record_function_or_nullcontext("Forward"),
self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
):
- # TODO(Peiqi): delete this
- model_kwargs_extra = {}
-
+ model_kwargs_extra = self._build_model_kwargs_extra(per_req_additional_information, num_scheduled_tokens_np)
model_output = self.model(
input_ids=input_ids,
positions=positions,
@@ -366,9 +152,15 @@ class GPUARModelRunner(OmniGPUModelRunner):
hidden_states = model_output
aux_hidden_states = None
- hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
- # TODO(Peiqi): move the merge stage out here
-
+ # hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+ # hidden_states
+ # )
+ multimodal_outputs = model_output.multimodal_outputs
+ hidden_states = model_output.text_hidden_states
+ # The model side may return per-request additional_information updates (model-agnostic channel).
+ # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
+ # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
+ self._process_additional_information_updates(multimodal_outputs)
if not self.broadcast_pp_output:
# Common case.
if not get_pp_group().is_last_rank:
@@ -484,12 +276,42 @@ class GPUARModelRunner(OmniGPUModelRunner):
# Convert to per-request tensors on CPU
hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
+ # pooler_output: list[torch.Tensor | None] = []
pooler_output: list[torch.Tensor | None] = []
prev_logits_index = 0
- for logits_index in logits_indices:
- pooler_output.append(hidden_states_cpu[prev_logits_index : logits_index + 1])
+ for rid, logits_index in zip(req_ids_output_copy, logits_indices):
+ # Base payload: hidden slice for this request in this iteration
+ hidden_slice = hidden_states_cpu[prev_logits_index : logits_index + 1]
+ payload: dict[str, object] = {"hidden": hidden_slice}
+ # Merge multimodal_outputs if present
+ if isinstance(multimodal_outputs, dict) and multimodal_outputs:
+ mm_payload: dict[str, object] = {}
+ for k, v in multimodal_outputs.items():
+ try:
+ # Case 1: tensor aligned on token dimension
+ if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
+ mm_payload[k] = v.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
+ # Case 2: nested dict of tensors aligned on token dimension (e.g., selected_hidden_layers)
+ elif isinstance(v, dict):
+ sub_dict: dict[str, torch.Tensor] = {}
+ for sk, sv in v.items():
+ if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
+ sub_dict[str(sk)] = (
+ sv.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
+ )
+ if sub_dict:
+ mm_payload[k] = sub_dict
+ elif isinstance(v, list):
+ element: torch.Tensor = v[0]
+ multimodal_outputs[k] = v[1:] if len(v) > 1 else v
+ mm_payload[k] = element
+ except Exception as e:
+ # Best-effort; skip malformed entries
+ logger.error(f"Error in merge multimodal outputs: {e}")
+ if mm_payload:
+ payload.update(mm_payload)
+ pooler_output.append(payload) # type: ignore[arg-type]
prev_logits_index = logits_index + 1
-
output = OmniModelRunnerOutput(
req_ids=req_ids_output_copy,
req_id_to_index=req_id_to_index_output_copy,
@@ -533,3 +355,5 @@ class GPUARModelRunner(OmniGPUModelRunner):
else:
merged[k] = v
setattr(req_state, "additional_information_cpu", merged)
+
+ # ===== Helper functions extracted for clarity and reuse =====
\ No newline at end of file
diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py
index 98826a93e..0c34e1ad6 100644
--- a/vllm_omni/worker/gpu_ar_worker.py
+++ b/vllm_omni/worker/gpu_ar_worker.py
@@ -13,7 +13,23 @@ from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
class GPUARWorker(GPUWorker):
+ """GPU worker for autoregressive omni model stages.
+
+ Extends the base GPUWorker to initialize and manage autoregressive
+ model runners for text generation stages (e.g., thinker stages).
+ """
+
def init_device(self):
+ """Initialize the GPU device and autoregressive model runner.
+
+ Sets up CUDA device, initializes distributed environment,
+ sets random seed, and creates a GPUARModelRunner instance.
+
+ Raises:
+ ValueError: If free GPU memory is insufficient for the
+ requested memory utilization
+ RuntimeError: If device type is not supported
+ """
if self.device_config.device.type == "cuda":
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
diff --git a/vllm_omni/worker/gpu_diffusion_model_runner.py b/vllm_omni/worker/gpu_diffusion_model_runner.py
index d381db8d1..401ff91bd 100644
--- a/vllm_omni/worker/gpu_diffusion_model_runner.py
+++ b/vllm_omni/worker/gpu_diffusion_model_runner.py
@@ -31,7 +31,7 @@ logger = logging.getLogger(__name__)
class GPUDiffusionModelRunner(OmniGPUModelRunner):
- """Diffusion model runner for vLLM-omni (non-autoregressive).
+ """Diffusion model runner for vLLM-Omni (non-autoregressive).
- Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
- Does not compute logits or perform token sampling.
diff --git a/vllm_omni/worker/gpu_diffusion_worker.py b/vllm_omni/worker/gpu_diffusion_worker.py
index 77f1f03fb..4731e26f1 100644
--- a/vllm_omni/worker/gpu_diffusion_worker.py
+++ b/vllm_omni/worker/gpu_diffusion_worker.py
@@ -13,7 +13,24 @@ from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
class GPUDiffusionWorker(GPUWorker):
+ """GPU worker for diffusion-based omni model stages.
+
+ Extends the base GPUWorker to initialize and manage diffusion model
+ runners for non-autoregressive generation stages (e.g., image or
+ audio generation stages).
+ """
+
def init_device(self):
+ """Initialize the GPU device and diffusion model runner.
+
+ Sets up CUDA device, initializes distributed environment,
+ sets random seed, and creates a GPUDiffusionModelRunner instance.
+
+ Raises:
+ ValueError: If free GPU memory is insufficient for the
+ requested memory utilization
+ RuntimeError: If device type is not supported
+ """
if self.device_config.device.type == "cuda":
# This env var set by Ray causes exceptions with graph building.
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py
new file mode 100644
index 000000000..d18b7e956
--- /dev/null
+++ b/vllm_omni/worker/gpu_generation_model_runner.py
@@ -0,0 +1,582 @@
+"""Code2Wav GPU Model Runner for vLLM-omni.
+
+Handles direct conversion from codec codes to audio waveforms for Qwen3 Omni MoE Code2Wav.
+This is a non-autoregressive model that doesn't require sampling or logits computation.
+"""
+
+from __future__ import annotations
+
+import gc
+import logging
+from typing import Any
+
+import numpy as np
+import torch
+from vllm.config import CUDAGraphMode
+from vllm.forward_context import BatchDescriptor
+from vllm.multimodal.inputs import MultiModalKwargs
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.spec_decode.eagle import EagleProposer
+from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker.gpu_model_runner import (
+ EMPTY_MODEL_RUNNER_OUTPUT,
+ AsyncGPUModelRunnerOutput,
+ IntermediateTensors,
+ PerLayerAttnMetadata,
+ get_pp_group,
+ set_forward_context,
+)
+from vllm.v1.worker.ubatch_utils import UBatchSlices
+from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
+
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+
+logger = logging.getLogger(__name__)
+
+
+class GPUGenerationModelRunner(OmniGPUModelRunner):
+ """Generation model runner for vLLM-omni (non-autoregressive).
+
+ - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
+ - Does not compute logits or perform token sampling.
+ - Executes generation process and returns tensors via `pooler_output`.
+ """
+
+ def _preprocess(
+ self,
+ scheduler_output: SchedulerOutput,
+ intermediate_tensors: IntermediateTensors | None = None,
+ ubatch_slices: UBatchSlices | None = None,
+ num_tokens_after_padding: torch.Tensor | None = None,
+ ) -> tuple[
+ int,
+ int,
+ torch.Tensor | None,
+ torch.Tensor | None,
+ torch.Tensor | None,
+ torch.Tensor,
+ IntermediateTensors | None,
+ dict[str, Any],
+ ]:
+ # Input token count for this iteration (not used by diffusion, but
+ # retained to keep DP padding/ordering consistent)
+ num_input_tokens = scheduler_output.total_num_scheduled_tokens
+ num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
+ num_input_tokens += num_pad
+
+ # _prepare_inputs may reorder the batch, so we must gather multi
+ # modal outputs after that to ensure the correct order
+ if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
+ # Run the multimodal encoder if any.
+ self._execute_mm_encoder(scheduler_output)
+ mm_embeds = self._gather_mm_embeddings(scheduler_output)
+
+ # NOTE(woosuk): To unify token ids and soft tokens (vision
+ # embeddings), we always use embeddings (rather than token ids)
+ # as input to the multimodal model, even when the input is text.
+ inputs_embeds_scheduled = self.model.get_input_embeddings(
+ input_ids=self.input_ids.gpu[:num_input_tokens],
+ multimodal_embeddings=mm_embeds or None,
+ )
+
+ # TODO(woosuk): Avoid the copy. Optimize.
+ self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled)
+
+ input_ids = self.input_ids.gpu[:num_input_tokens]
+ inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+ model_kwargs = {
+ **self._init_model_kwargs(num_input_tokens),
+ **self._extract_mm_kwargs(scheduler_output),
+ }
+ elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
+ # Get the input embeddings for the tokens that are not input embeds,
+ # then put them into the appropriate positions.
+ # TODO(qthequartermasterman): Since even when prompt embeds are
+ # enabled, (a) not all requests will use prompt embeds, and (b)
+ # after the initial prompt is processed, the rest of the generated
+ # tokens will be token ids, it is not desirable to have the
+ # embedding layer outside of the CUDA graph all the time. The v0
+ # engine avoids this by "double compiling" the CUDA graph, once
+ # with input_ids and again with inputs_embeds, for all num_tokens.
+ # If a batch only has token ids, then including the embedding layer
+ # in the CUDA graph will be more performant (like in the else case
+ # below).
+ token_ids_idx = self.is_token_ids.gpu[:num_input_tokens].nonzero(as_tuple=False).squeeze(1)
+ # Some tokens ids may need to become embeds
+ if token_ids_idx.numel() > 0:
+ token_ids = self.input_ids.gpu[token_ids_idx]
+ tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
+ self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
+
+ inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+ model_kwargs = self._init_model_kwargs(num_input_tokens)
+ input_ids = self.input_ids.gpu[:num_input_tokens]
+ else:
+ # For text-only models, we use token ids as input.
+ # While it is possible to use embeddings as input just like the
+ # multimodal models, it is not desirable for performance since
+ # then the embedding layer is not included in the CUDA graph.
+ input_ids = self.input_ids.gpu[:num_input_tokens]
+ inputs_embeds = None
+ model_kwargs = self._init_model_kwargs(num_input_tokens)
+ if self.uses_mrope:
+ positions = self.mrope_positions.gpu[:, :num_input_tokens]
+ else:
+ positions = self.positions.gpu[:num_input_tokens]
+
+ if get_pp_group().is_first_rank:
+ intermediate_tensors = None
+ else:
+ intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+ num_input_tokens, intermediate_tensors, True
+ )
+
+ if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
+ encoder_inputs = self._extract_encoder_inputs(scheduler_output)
+ model_kwargs.update(encoder_inputs)
+
+ return (
+ num_input_tokens,
+ num_input_tokens,
+ num_tokens_after_padding,
+ input_ids,
+ inputs_embeds,
+ positions,
+ intermediate_tensors,
+ model_kwargs,
+ )
+
+ @torch.inference_mode()
+ def execute_model(
+ self,
+ scheduler_output: SchedulerOutput,
+ intermediate_tensors: IntermediateTensors | None = None,
+ ) -> OmniModelRunnerOutput | IntermediateTensors:
+ with record_function_or_nullcontext("Preprocess"):
+ with self.synchronize_input_prep():
+ # Update persistent batch states.
+ self._update_states(scheduler_output)
+ if not scheduler_output.total_num_scheduled_tokens:
+ return EMPTY_MODEL_RUNNER_OUTPUT
+
+ # Prepare the decoder inputs.
+ (
+ attn_metadata,
+ logits_indices,
+ spec_decode_metadata,
+ num_scheduled_tokens_np,
+ spec_decode_common_attn_metadata,
+ max_query_len,
+ ubatch_slices,
+ num_tokens_after_padding,
+ ) = self._prepare_inputs(scheduler_output)
+
+ # Input token count for this iteration (not used by diffusion, but
+ # retained to keep DP padding/ordering consistent)
+ (
+ num_scheduled_tokens,
+ num_input_tokens,
+ num_tokens_across_dp,
+ input_ids,
+ inputs_embeds,
+ positions,
+ intermediate_tensors,
+ model_kwargs,
+ ) = self._preprocess(
+ scheduler_output,
+ intermediate_tensors,
+ ubatch_slices,
+ num_tokens_after_padding,
+ )
+
+ cudagraph_runtime_mode = CUDAGraphMode.NONE
+ # Run the model.
+ # Use persistent buffers for CUDA graphs.
+ with (
+ set_forward_context(
+ attn_metadata,
+ self.vllm_config,
+ num_tokens=num_input_tokens,
+ num_tokens_across_dp=num_tokens_across_dp,
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
+ batch_descriptor=None,
+ ubatch_slices=ubatch_slices,
+ ),
+ record_function_or_nullcontext("Forward"),
+ self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+ ):
+ outputs = self._run_generation_model(
+ input_ids=input_ids,
+ positions=positions,
+ intermediate_tensors=intermediate_tensors,
+ inputs_embeds=inputs_embeds,
+ multimodal_kwargs=model_kwargs,
+ logits_indices=logits_indices,
+ )
+
+ _, multimodal_outputs = self.extract_multimodal_outputs(outputs)
+ # Ensure one tensor per request, map to CPU for output struct
+ pooler_output: list[torch.Tensor | None] = []
+ if isinstance(multimodal_outputs, torch.Tensor):
+ # If model returned a single stacked tensor, split by requests
+ assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
+ for i in range(self.input_batch.num_reqs):
+ pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous())
+ elif isinstance(multimodal_outputs, list):
+ for out in multimodal_outputs:
+ pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+ elif isinstance(multimodal_outputs, dict):
+ for out in multimodal_outputs.values():
+ pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+ else:
+ raise RuntimeError("Unsupported diffusion output type")
+
+ output = OmniModelRunnerOutput(
+ req_ids=self.input_batch.req_ids,
+ req_id_to_index=self.input_batch.req_id_to_index,
+ sampled_token_ids=[],
+ logprobs=None,
+ prompt_logprobs_dict={},
+ pooler_output=pooler_output,
+ kv_connector_output=kv_connector_output,
+ num_nans_in_logits={},
+ )
+
+ if not self.use_async_scheduling:
+ return output
+
+ return AsyncGPUModelRunnerOutput(
+ model_runner_output=output,
+ sampled_token_ids=[],
+ invalid_req_indices=[],
+ async_output_copy_stream=self.async_output_copy_stream,
+ )
+
+ def _run_generation_model(
+ self,
+ *,
+ input_ids: torch.Tensor,
+ positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None,
+ inputs_embeds: torch.Tensor | None,
+ multimodal_kwargs: dict,
+ logits_indices: torch.Tensor,
+ ) -> torch.Tensor | list[torch.Tensor]:
+ """Run generation from codec codes to waveforms.
+
+ Args:
+ scheduler_output: Contains codec codes in input_ids or additional info
+ intermediate_tensors: PP intermediate tensors if applicable
+
+ Returns:
+ Audio waveforms: [batch, 1, waveform_len] or list of tensors
+ """
+ # Keep inputs identical to AR runner
+ kwargs = dict(
+ input_ids=input_ids,
+ positions=positions,
+ intermediate_tensors=intermediate_tensors,
+ inputs_embeds=inputs_embeds,
+ **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
+ sampling_metadata=self.input_batch.sampling_metadata,
+ logits_index=logits_indices,
+ sampler=self.sampler,
+ )
+
+ if hasattr(self.model, "forward"):
+ return self.model.forward(**kwargs)
+
+ raise RuntimeError(
+ "The loaded model does not expose diffusion interfaces 'sample', "
+ "'forward', or 'diffuse'. Please implement one of them or adapt the runner."
+ )
+
+ @torch.inference_mode()
+ def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+ logger.warning("Dummy sampler run is not implemented for generation model")
+ return None
+
+ @torch.inference_mode()
+ def _dummy_run(
+ self,
+ num_tokens: int,
+ cudagraph_runtime_mode: CUDAGraphMode | None = None,
+ force_attention: bool = False,
+ uniform_decode: bool = False,
+ skip_eplb: bool = False,
+ is_profile: bool = False,
+ create_mixed_batch: bool = False,
+ remove_lora: bool = True,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Run a dummy forward pass to warm up/profile run or capture the
+ CUDA graph for the model.
+
+ Args:
+ num_tokens: Number of tokens to run the dummy forward pass.
+ cudagraph_runtime_mode: used to control the behavior.
+ - if not set will determine the cudagraph mode based on using
+ the self.cudagraph_dispatcher.
+ - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
+ - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
+ - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
+ needed.
+ force_attention: If True, always create attention metadata. Used to
+ warm up attention backend when mode is NONE.
+ uniform_decode: If True, the batch is a uniform decode batch.
+ skip_eplb: If True, skip EPLB state update.
+ is_profile: If True, this is a profile run.
+ create_mixed_batch: If True, create a mixed batch with both decode
+ (1 token) and prefill (multiple tokens) requests.
+ remove_lora: If False, dummy LoRAs are not destroyed after the run
+ """
+ assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
+ CUDAGraphMode.NONE,
+ CUDAGraphMode.PIECEWISE,
+ CUDAGraphMode.FULL,
+ }
+
+ # If cudagraph_mode.decode_mode() == FULL and
+ # cudagraph_mode.separate_routine(). This means that we are using
+ # different graphs and/or modes for mixed prefill-decode batches vs.
+ # uniform decode batches. A uniform decode batch means that all
+ # requests have identical query length, except a potential virtual
+ # request (shorter) in the batch account for padding.
+ # Uniform decode batch could either be common pure decode, where
+ # max_query_len == 1, or speculative decode, where
+ # max_query_len == 1 + num_spec_decode_tokens.
+
+ # When setting max_query_len = 1, we switch to and capture the optimized
+ # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
+ # for GQA/MQA.
+ max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
+
+ # Set num_scheduled_tokens based on num_tokens and max_num_seqs
+ # for dummy run with LoRA so that the num_reqs collectively
+ # has num_tokens in total.
+ assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+ max_num_reqs = self.scheduler_config.max_num_seqs
+
+ num_reqs = min(num_tokens, max_num_reqs)
+ min_tokens_per_req = num_tokens // num_reqs
+ num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
+ num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+
+ assert sum(num_scheduled_tokens_list) == num_tokens
+ assert len(num_scheduled_tokens_list) == num_reqs
+ num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
+
+ num_tokens_after_padding = None
+
+ # If we failed to microbatch, currently need to resynchronize
+ # TODO(lucas,sage): we should be able to avoid this second sync by
+ # refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
+ # a single `coordinate_batch_across_dp` function.
+ if num_tokens_after_padding is None:
+ num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
+ num_tokens_after_padding = num_tokens + num_pad
+ else:
+ num_tokens_across_dp = num_tokens_after_padding
+ num_tokens_after_padding = int(num_tokens_after_padding[0].item())
+
+ attn_metadata: PerLayerAttnMetadata | None = None
+
+ # If force_attention is True, we always capture attention. Otherwise,
+ # it only happens for cudagraph_runtime_mode=FULL.
+ if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
+ attn_metadata = {}
+
+ seq_lens = max_query_len
+ self.seq_lens.np[:num_reqs] = seq_lens
+ self.seq_lens.np[num_reqs:] = 0
+ self.seq_lens.copy_to_gpu()
+
+ cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+ self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+ self.query_start_loc.copy_to_gpu()
+
+ for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups):
+ common_attn_metadata = CommonAttentionMetadata(
+ query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
+ query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
+ seq_lens=self.seq_lens.gpu[:num_reqs],
+ seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
+ num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs], # noqa: E501
+ num_reqs=num_reqs,
+ num_actual_tokens=num_tokens,
+ max_query_len=max_query_len,
+ max_seq_len=self.max_model_len,
+ block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs),
+ slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens],
+ causal=True,
+ )
+ for attn_group in self.attn_groups[kv_cache_group_id]:
+ assert type(attn_metadata) is dict
+ attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(
+ common_attn_metadata
+ ) # noqa: E501
+ for layer_name in attn_group.layer_names:
+ attn_metadata[layer_name] = attn_metadata_i
+
+ with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora):
+ model_kwargs = self._init_model_kwargs(num_tokens)
+ if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
+ input_ids = None
+ inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+ model_kwargs = {
+ **model_kwargs,
+ **self._dummy_mm_kwargs(num_reqs),
+ }
+ elif self.enable_prompt_embeds:
+ input_ids = None
+ inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+ model_kwargs = self._init_model_kwargs(num_tokens)
+ else:
+ input_ids = self.input_ids.gpu[:num_tokens]
+ inputs_embeds = None
+
+ if self.uses_mrope:
+ positions = self.mrope_positions.gpu[:, :num_tokens]
+ else:
+ positions = self.positions.gpu[:num_tokens]
+
+ if get_pp_group().is_first_rank:
+ intermediate_tensors = None
+ else:
+ if self.intermediate_tensors is None:
+ self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
+ batch_size=self.max_num_tokens,
+ dtype=self.model_config.dtype,
+ device=self.device,
+ )
+
+ intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens, None, False)
+
+ # filter out the valid batch descriptor
+ _cg_mode, batch_descriptor = (
+ self.cudagraph_dispatcher.dispatch(
+ BatchDescriptor(
+ num_tokens=num_tokens_after_padding,
+ uniform_decode=uniform_decode,
+ )
+ )
+ if not is_profile
+ else (CUDAGraphMode.NONE, None)
+ )
+ if cudagraph_runtime_mode is not None:
+ # we allow forcing NONE when the dispatcher disagrees to support
+ # warm ups for cudagraph capture
+ assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, (
+ f"Cudagraph runtime mode mismatch at dummy_run. "
+ f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
+ )
+ else:
+ cudagraph_runtime_mode = _cg_mode
+
+ with (
+ self.maybe_randomize_inputs(input_ids),
+ set_forward_context(
+ attn_metadata,
+ self.vllm_config,
+ num_tokens=num_tokens_after_padding,
+ num_tokens_across_dp=num_tokens_across_dp,
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
+ batch_descriptor=batch_descriptor,
+ ),
+ ):
+ outputs = self.model(
+ input_ids=input_ids,
+ positions=positions,
+ intermediate_tensors=intermediate_tensors,
+ inputs_embeds=inputs_embeds,
+ **model_kwargs,
+ )
+
+ if self.use_aux_hidden_state_outputs:
+ hidden_states, _ = outputs
+ else:
+ hidden_states = outputs
+
+ if self.speculative_config and self.speculative_config.use_eagle():
+ assert isinstance(self.drafter, EagleProposer)
+ self.drafter.dummy_run(num_tokens)
+
+ # This is necessary to avoid blocking DP.
+ # For dummy runs, we typically skip EPLB since we don't have any real
+ # requests to process.
+ # However, in DP settings, there may be cases when some DP ranks do
+ # not have any requests to process, so they're executing dummy batches.
+ # In such cases, we still have to trigger EPLB to make sure
+ # ranks execute the rearrangement in synchronization.
+
+ # logit_indices = np.cumsum(num_scheduled_tokens) - 1
+ hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
+ return hidden_states, None
+
+ def profile_run(self) -> None:
+ # Profile with multimodal encoder & encoder cache.
+ if self.supports_mm_inputs:
+ if self.model_config.multimodal_config.skip_mm_profiling:
+ logger.info("Skipping memory profiling for multimodal encoder and encoder cache.")
+ else:
+ mm_budget = self.mm_budget
+ assert mm_budget is not None
+
+ if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
+ # NOTE: Currently model is profiled with a single non-text
+ # modality with the max possible input tokens even when
+ # it supports multiple.
+ dummy_modality = mm_budget.get_modality_with_max_tokens()
+ max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[dummy_modality]
+
+ logger.info(
+ "Encoder cache will be initialized with a budget "
+ "of %s tokens, and profiled with %s %s items of the maximum feature size.",
+ encoder_budget,
+ max_mm_items_per_batch,
+ dummy_modality,
+ )
+
+ # Create dummy batch of multimodal inputs.
+ batched_dummy_mm_inputs = self._get_mm_dummy_batch(
+ dummy_modality,
+ max_mm_items_per_batch,
+ )
+
+ # Run multimodal encoder.
+ dummy_encoder_outputs = self.model.get_multimodal_embeddings(**batched_dummy_mm_inputs)
+
+ sanity_check_mm_encoder_outputs(
+ dummy_encoder_outputs,
+ expected_num_items=max_mm_items_per_batch,
+ )
+
+ # NOTE: This happens when encoder cache needs to store
+ # the embeddings that encoder outputs are scattered onto.
+ # In this case we create dummy embeddings of size
+ # (encode_budget, hidden_size) and scatter encoder
+ # output into it.
+ encoder_output_shape = dummy_encoder_outputs[0].shape
+ if encoder_output_shape[0] < encoder_budget:
+ expanded_outputs = []
+ for output in dummy_encoder_outputs:
+ expanded = output.new_zeros((encoder_budget, encoder_output_shape[-1]))
+ num_tokens = output.shape[0]
+ expanded[:num_tokens].copy_(output)
+ expanded_outputs.append(expanded)
+
+ dummy_encoder_outputs = expanded_outputs
+
+ # Cache the dummy encoder outputs.
+ self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
+
+ # Add `is_profile` here to pre-allocate communication buffers
+ hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
+ if get_pp_group().is_last_rank:
+ pass
+ self._sync_device()
+ del hidden_states
+ self.encoder_cache.clear()
+ gc.collect()
diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py
new file mode 100644
index 000000000..49dc10c7b
--- /dev/null
+++ b/vllm_omni/worker/gpu_generation_worker.py
@@ -0,0 +1,84 @@
+import gc
+import os
+
+import torch
+from vllm.model_executor import set_random_seed
+from vllm.platforms import current_platform
+from vllm.utils import GiB_bytes, MemorySnapshot
+from vllm.v1.utils import report_usage_stats
+from vllm.v1.worker.gpu_worker import Worker as GPUWorker
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+
+from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
+
+
+class GPUGenerationWorker(GPUWorker):
+ """GPU Worker for Generation model (non-autoregressive waveform generation).
+
+ Usage in stage config:
+ worker_cls: "vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner"
+ """
+
+ def init_device(self):
+ """Initialize CUDA device and distributed environment."""
+ if self.device_config.device.type == "cuda":
+ # torch.distributed.all_reduce does not free the input tensor until
+ # the synchronization point. This causes the memory usage to grow
+ # as the number of all_reduce calls increases. This env var disables
+ # this behavior.
+ # Related issue:
+ # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
+ os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
+
+ # This env var set by Ray causes exceptions with graph building.
+ os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+ self.device = torch.device(f"cuda:{self.local_rank}")
+ current_platform.set_device(self.device)
+
+ current_platform.check_if_supports_dtype(self.model_config.dtype)
+
+ # Initialize the distributed environment BEFORE taking
+ # memory snapshot
+ # This ensures NCCL buffers are allocated before we measure
+ # available memory
+ init_worker_distributed_environment(
+ self.vllm_config,
+ self.rank,
+ self.distributed_init_method,
+ self.local_rank,
+ current_platform.dist_backend,
+ )
+
+ # Set random seed.
+ set_random_seed(self.model_config.seed)
+
+ # Now take memory snapshot after NCCL is initialized
+ gc.collect()
+ torch.cuda.empty_cache()
+
+ # take current memory snapshot
+ self.init_snapshot = MemorySnapshot()
+ self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization
+ if self.init_snapshot.free_memory < self.requested_memory:
+
+ def GiB(b):
+ return round(b / GiB_bytes, 2)
+
+ raise ValueError(
+ f"Free memory on device "
+ f"({GiB(self.init_snapshot.free_memory)}/"
+ f"{GiB(self.init_snapshot.total_memory)} GiB) on startup "
+ f"is less than desired GPU memory utilization "
+ f"({self.cache_config.gpu_memory_utilization}, "
+ f"{GiB(self.requested_memory)} GiB). Decrease GPU memory "
+ f"utilization or reduce GPU memory used by other processes."
+ )
+ else:
+ raise RuntimeError(f"Not support device type: {self.device_config.device}")
+
+ # Construct the model runner
+ self.model_runner: GPUGenerationModelRunner = GPUGenerationModelRunner(self.vllm_config, self.device)
+
+ if self.rank == 0:
+ # If usage stat is enabled, collect relevant info.
+ report_usage_stats(self.vllm_config)
diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py
index 55e5e1391..f2d9c1111 100644
--- a/vllm_omni/worker/gpu_model_runner.py
+++ b/vllm_omni/worker/gpu_model_runner.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
import numpy as np
import torch
@@ -6,15 +6,18 @@ from vllm.config import CUDAGraphMode
from vllm.distributed.parallel_state import get_pp_group
from vllm.forward_context import BatchDescriptor, set_forward_context
from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import supports_mrope
from vllm.model_executor.models.interfaces_base import VllmModelForPooling
from vllm.sampling_params import SamplingType
from vllm.utils import LazyLoader, cdiv
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata, split_attn_metadata
+from vllm.v1.attention.backends.utils import (
+ CommonAttentionMetadata,
+ split_attn_metadata,
+)
from vllm.v1.spec_decode.eagle import EagleProposer
from vllm.v1.worker.gpu_input_batch import CachedRequestState
-from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata
from vllm.v1.worker.ubatch_splitting import ubatch_split
+from vllm.v1.worker.ubatch_utils import UBatchSlices
if TYPE_CHECKING:
from vllm.v1.core.sched.output import SchedulerOutput
@@ -30,42 +33,6 @@ logger = init_logger(__name__)
class OmniGPUModelRunner(GPUModelRunner):
- def _init_mrope_positions(self, req_state: CachedRequestState):
- image_grid_thw = []
- video_grid_thw = []
- second_per_grid_ts = []
- audio_feature_lengths = []
- use_audio_in_video = False
- for mm_feature in req_state.mm_features:
- mm_item = mm_feature.data
- if mm_item is None:
- continue
- mm_input = mm_item.get_data()
- if (t := mm_input.get("image_grid_thw")) is not None:
- image_grid_thw.append(t.tolist())
- if (t := mm_input.get("video_grid_thw")) is not None:
- video_grid_thw.append(t.tolist())
- if (t := mm_input.get("second_per_grid_ts")) is not None:
- second_per_grid_ts.append(t)
- if (t := mm_input.get("audio_feature_lengths")) is not None:
- audio_feature_lengths.append(t)
- # Check for use_audio_in_video
- use_audio_in_video_value = mm_input.get("use_audio_in_video")
- if use_audio_in_video_value is not None:
- use_audio_in_video = bool(use_audio_in_video_value.item())
-
- assert supports_mrope(self.get_model()), "M-RoPE support is not implemented."
-
- req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions(
- req_state.prompt_token_ids,
- hf_config=self.model_config.hf_config,
- image_grid_thw=image_grid_thw,
- video_grid_thw=video_grid_thw,
- second_per_grid_ts=second_per_grid_ts,
- audio_feature_lengths=audio_feature_lengths,
- use_audio_in_video=use_audio_in_video,
- )
-
def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
"""Update the cached states and the persistent batch with the scheduler
output.
@@ -581,3 +548,291 @@ class OmniGPUModelRunner(GPUModelRunner):
logit_indices = np.cumsum(num_scheduled_tokens) - 1
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
return hidden_states, hidden_states[logit_indices]
+
+ def _preprocess(
+ self,
+ scheduler_output: "SchedulerOutput",
+ num_scheduled_tokens_np: np.ndarray,
+ intermediate_tensors: Optional[IntermediateTensors] = None,
+ ubatch_slices: Optional[UBatchSlices] = None,
+ num_tokens_after_padding: Optional[torch.Tensor] = None,
+ ) -> tuple[
+ int,
+ int,
+ Optional[torch.Tensor],
+ Optional[torch.Tensor],
+ Optional[torch.Tensor],
+ torch.Tensor,
+ Optional[IntermediateTensors],
+ dict[str, Any],
+ Optional[dict[str, dict]],
+ ]:
+ num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+ if ubatch_slices:
+ assert num_tokens_after_padding is not None
+ num_input_tokens = int(num_tokens_after_padding[0].item() * 2)
+ self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
+ elif ubatch_slices is None:
+ num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)
+ num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
+ num_input_tokens += num_pad
+
+ # _prepare_inputs may reorder the batch, so we must gather multi
+ # modal outputs after that to ensure the correct order
+ per_req_additional_information: Optional[dict[str, dict]] = None
+ if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
+ # Build multimodal inputs and overlay prompt embeds; collect per-request info
+ per_req_additional_information = self._build_mm_inputs_and_overlays(
+ scheduler_output, num_scheduled_tokens, num_scheduled_tokens_np
+ )
+
+ input_ids = self.input_ids.gpu[:num_input_tokens]
+ inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+ model_kwargs = {
+ **self._init_model_kwargs(num_scheduled_tokens),
+ **self._extract_mm_kwargs(scheduler_output),
+ }
+ elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
+ # Get the input embeddings for the tokens that are not input embeds,
+ # then put them into the appropriate positions.
+ # TODO(qthequartermasterman): Since even when prompt embeds are
+ # enabled, (a) not all requests will use prompt embeds, and (b)
+ # after the initial prompt is processed, the rest of the generated
+ # tokens will be token ids, it is not desirable to have the
+ # embedding layer outside of the CUDA graph all the time. The v0
+ # engine avoids this by "double compiling" the CUDA graph, once
+ # with input_ids and again with inputs_embeds, for all num_tokens.
+ # If a batch only has token ids, then including the embedding layer
+ # in the CUDA graph will be more performant (like in the else case
+ # below).
+ token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1)
+ # Some tokens ids may need to become embeds
+ if token_ids_idx.numel() > 0:
+ token_ids = self.input_ids.gpu[token_ids_idx]
+ tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
+ self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
+
+ inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+ model_kwargs = self._init_model_kwargs(num_input_tokens)
+ input_ids = None
+ else:
+ # For text-only models, we use token ids as input.
+ # While it is possible to use embeddings as input just like the
+ # multimodal models, it is not desirable for performance since
+ # then the embedding layer is not included in the CUDA graph.
+ input_ids = self.input_ids.gpu[:num_input_tokens]
+ inputs_embeds = None
+ model_kwargs = self._init_model_kwargs(num_input_tokens)
+ if self.uses_mrope:
+ positions = self.mrope_positions.gpu[:, :num_input_tokens]
+ else:
+ positions = self.positions.gpu[:num_input_tokens]
+
+ if get_pp_group().is_first_rank:
+ intermediate_tensors = None
+ else:
+ intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+ num_input_tokens, intermediate_tensors, True
+ )
+
+ if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
+ encoder_inputs = self._extract_encoder_inputs(scheduler_output)
+ model_kwargs.update(encoder_inputs)
+
+ return (
+ num_scheduled_tokens,
+ num_input_tokens,
+ num_tokens_after_padding,
+ input_ids,
+ inputs_embeds,
+ positions,
+ intermediate_tensors,
+ model_kwargs,
+ per_req_additional_information,
+ )
+
+ def _decode_and_store_request_payloads(self, scheduler_output: "SchedulerOutput") -> None:
+ """Decode per-request prompt_embeds and additional_information for newly
+ scheduled requests and store them to CPU in the request state.
+ This version avoids hard dependency on payload classes by duck-typing."""
+ try:
+ new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
+ if not new_reqs:
+ return
+ for nr in new_reqs:
+ req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
+ if req_id is None:
+ continue
+ # prompt_embeds
+ payload_pe = getattr(nr, "prompt_embeds", None)
+ pe_cpu = None
+ if payload_pe is not None:
+ if isinstance(payload_pe, torch.Tensor):
+ pe_cpu = payload_pe.detach().to("cpu").contiguous()
+ else:
+ # Try duck-typing a payload with data/shape/dtype
+ data = getattr(payload_pe, "data", None)
+ shape = getattr(payload_pe, "shape", None)
+ if data is not None and shape is not None:
+ dt = np.dtype(getattr(payload_pe, "dtype", "float32"))
+ arr = np.frombuffer(data, dtype=dt)
+ arr = arr.reshape(shape)
+ pe_cpu = torch.from_numpy(arr)
+ if pe_cpu is not None and req_id in self.requests:
+ setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
+ # additional_information
+ payload_info = getattr(nr, "additional_information", None)
+ if payload_info is not None:
+ info_dict = {}
+ if isinstance(payload_info, dict):
+ info_dict = payload_info
+ else:
+ # Try duck-typing a payload with entries, each entry may have
+ # tensor_data/tensor_dtype/tensor_shape or list_data
+ entries = getattr(payload_info, "entries", None)
+ if isinstance(entries, dict):
+ for k, entry in entries.items():
+ tensor_data = getattr(entry, "tensor_data", None)
+ if tensor_data is not None:
+ dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
+ arr = np.frombuffer(tensor_data, dtype=dt)
+ arr = arr.reshape(getattr(entry, "tensor_shape", ()))
+ info_dict[k] = torch.from_numpy(arr)
+ else:
+ info_dict[k] = getattr(entry, "list_data", None)
+ if info_dict and req_id in self.requests:
+ setattr(self.requests[req_id], "additional_information_cpu", info_dict)
+ except Exception as e:
+ logger.error(f"Error decoding prompt_embeds / additional_information: {e}")
+
+ def _gather_runtime_additional_information(self) -> list[dict]:
+ """Gather per-request additional_information stored in request state in batch order."""
+ per_req_runtime_info = []
+ for req_id in self.input_batch.req_ids:
+ req_state = self.requests.get(req_id)
+ info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
+ per_req_runtime_info.append(info if isinstance(info, dict) else {})
+ return per_req_runtime_info
+
+ def _compute_request_token_spans(self, num_scheduled_tokens_np) -> list[tuple[int, int]]:
+ """Compute (start, end) token spans for each request within the flattened step sequence."""
+ req_token_spans: list[tuple[int, int]] = []
+ for req_index in range(len(self.input_batch.req_ids)):
+ start_offset = int(self.query_start_loc.cpu[req_index])
+ sched_tokens = int(num_scheduled_tokens_np[req_index])
+ req_token_spans.append((start_offset, start_offset + sched_tokens))
+ return req_token_spans
+
+ def _build_model_kwargs_extra(
+ self,
+ per_req_additional_information: Optional[dict[str, dict]],
+ num_scheduled_tokens_np,
+ ) -> dict:
+ """Build extra keyword arguments passed to the model for this step, including:
+ - additional_information_by_req_id: per-request additional information provided by preprocess for this step
+ - runtime_additional_information: per-request additional information stored in request state
+ - request_ids: the request id list in batch order
+ - request_token_spans: token spans per request within the flattened sequence
+ """
+ model_kwargs_extra: dict[str, object] = {}
+ if per_req_additional_information:
+ model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
+ try:
+ model_kwargs_extra["runtime_additional_information"] = self._gather_runtime_additional_information()
+ model_kwargs_extra["request_ids"] = self.input_batch.req_ids
+ model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(num_scheduled_tokens_np)
+ except Exception:
+ pass
+ return model_kwargs_extra
+
+ def _process_additional_information_updates(self, multimodal_outputs: object) -> None:
+ """Process model-provided per-request additional_information updates and merge into request state."""
+ try:
+ if not isinstance(multimodal_outputs, dict):
+ return
+ if (
+ "additional_information_update" not in multimodal_outputs
+ and "additional_information_update_by_req_id" not in multimodal_outputs
+ ):
+ return
+ updates_list = multimodal_outputs.get("additional_information_update")
+ if isinstance(updates_list, list):
+ for idx, upd in enumerate(updates_list):
+ if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
+ continue
+ req_id = self.input_batch.req_ids[idx]
+ self._merge_additional_information_update(req_id, upd)
+ updates_map = multimodal_outputs.get("additional_information_update_by_req_id")
+ if isinstance(updates_map, dict):
+ for req_id, upd in updates_map.items():
+ if not isinstance(upd, dict):
+ continue
+ if req_id not in self.requests:
+ continue
+ self._merge_additional_information_update(req_id, upd)
+ except Exception as e:
+ logger.error(
+ f"Error merging for requests:{self.input_batch.req_ids} "
+ f"additional information update: {e}, with the multimodal_outputs "
+ f"as {multimodal_outputs}"
+ )
+
+ def _build_mm_inputs_and_overlays(
+ self,
+ scheduler_output: "SchedulerOutput",
+ num_scheduled_tokens: int,
+ num_scheduled_tokens_np: np.ndarray,
+ ) -> dict[str, dict]:
+ """Run MM encoder, fill scheduled inputs_embeds, and overlay prompt embeds;
+ collect per-request additional_information for the prefill portion."""
+ # Run the multimodal encoder if any.
+ self._execute_mm_encoder(scheduler_output)
+ mm_embeds = self._gather_mm_embeddings(scheduler_output)
+ # Always use embeddings as inputs for MM models
+ inputs_embeds_scheduled = self.model.get_input_embeddings(
+ input_ids=self.input_ids.gpu[:num_scheduled_tokens],
+ multimodal_embeddings=mm_embeds or None,
+ )
+ self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
+ # Reset per-step additional information collector (deprecated concat path)
+ if hasattr(self, "_forward_additional_information"):
+ self._forward_additional_information = None
+ # Overlay and collect per-request info
+ return self._collect_additional_information_for_prefill(num_scheduled_tokens_np)
+
+ def _collect_additional_information_for_prefill(
+ self,
+ num_scheduled_tokens_np: np.ndarray,
+ ) -> dict[str, dict]:
+ """Overlay per-request prompt_embeds for the prefill portion and collect
+ additional_information slices for this step. Returns a map req_id -> dict."""
+ per_req_additional_information: dict[str, dict] = {}
+ for req_index, req_id in enumerate(self.input_batch.req_ids):
+ req_state = self.requests[req_id]
+ pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
+ addi_cpu = getattr(req_state, "additional_information_cpu", None)
+ num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index])
+ prompt_len = len(req_state.prompt_token_ids)
+ prompt_remaining = max(0, prompt_len - num_computed_tokens)
+ sched_tokens = int(num_scheduled_tokens_np[req_index])
+ overlay_len = min(sched_tokens, prompt_remaining)
+ if overlay_len <= 0 and not (isinstance(addi_cpu, dict) and addi_cpu):
+ continue
+ if overlay_len > 0 and pe_cpu is not None:
+ src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to(
+ dtype=self.dtype, device=self.device, non_blocking=True
+ )
+ start_offset = int(self.query_start_loc.cpu[req_index])
+ self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
+ if addi_cpu is not None and isinstance(addi_cpu, dict):
+ req_info: dict[str, object] = {}
+ for k, v in addi_cpu.items():
+ if isinstance(v, torch.Tensor):
+ req_info[k] = v.detach().to("cpu").contiguous()
+ elif isinstance(v, list):
+ req_info[k] = v
+ else:
+ req_info[k] = v
+ if req_info:
+ per_req_additional_information[req_id] = req_info
+ return per_req_additional_information
+ 
+