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..c1d840081 100644
--- a/README.md
+++ b/README.md
@@ -55,77 +55,18 @@ vLLM seamlessly supports most popular open-source models on HuggingFace, includi
## 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..157ec09a1
--- /dev/null
+++ b/docs/.nav.yml
@@ -0,0 +1,38 @@
+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: user_guide/examples/online_serving/README.md
+ - General:
+ - usage/*
+ - Configuration:
+ - configuration/README.md
+ - configuration/*
+ - Models:
+ - models/supported_models.md
+ - Features: features
+ - Developer Guide:
+ - contributing/README.md
+ - General:
+ - glob: contributing/*
+ flatten_single_child_sections: true
+ - Model Implementation:
+ - contributing/model/README.md
+ - CI: contributing/ci
+ - Design Documents: design
+ - Docs Guide: contributing/DOCS_GUIDE.md
+ - API Reference:
+ - api/README.md
+ - api/vllm_omni
+ - CLI Reference: cli
+ - Community:
+ - community/*
+ - Blog: https://blog.vllm.ai
+ - Forum: https://discuss.vllm.ai
+ - Slack: https://slack.vllm.ai
diff --git a/docs/README.md b/docs/README.md
new file mode 100644
index 000000000..f3e0af3b5
--- /dev/null
+++ b/docs/README.md
@@ -0,0 +1,60 @@
+---
+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 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:
+
+- 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
+
+vLLM-omni is flexible and easy to use with:
+
+- 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-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)
+
+For more information, checkout the following:
diff --git a/docs/api/README.md b/docs/api/README.md
new file mode 100644
index 000000000..9398fea11
--- /dev/null
+++ b/docs/api/README.md
@@ -0,0 +1,88 @@
+# 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_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.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.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..d1a0fbccf
--- /dev/null
+++ b/docs/community/contact_us.md
@@ -0,0 +1,6 @@
+# Contact Us
+
+- For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
+- For discussing with fellow users, please use the [vLLM Forum](https://discuss.vllm.ai/)
+- For coordinating contributions and development, please use [Slack](https://slack.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..48ceac159
--- /dev/null
+++ b/docs/configuration/stage_configs.md
@@ -0,0 +1,267 @@
+# Stage configs for vLLM-Omni
+
+In vLLM-Omni, it divides models in vLLM-omni into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other kinds of engines. According to different types of stages, such as AR stage or DiT stage, each stage can choose corresponding schedulers, model workers to load into Engines with a plug-in way.
+
+Therefore, the stage configs for a model have several main functions:
+
+- Claim partition of stages and corresponding class implementation in the folder 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
+```
+
+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..e4f6b8e7b
--- /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: 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
diff --git a/docs/DOCS_GUIDE.md b/docs/contributing/DOCS_GUIDE.md
similarity index 100%
rename from docs/DOCS_GUIDE.md
rename to docs/contributing/DOCS_GUIDE.md
diff --git a/docs/contributing/README.md b/docs/contributing/README.md
index 43bccdc5b..c27b7f32c 100644
--- a/docs/contributing/README.md
+++ b/docs/contributing/README.md
@@ -4,7 +4,7 @@ Thank you for your interest in contributing to vLLM-omni! This document provides
## Getting Started
-### Development Environment
+### Development Environment for vLLM
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:
@@ -19,8 +19,24 @@ Install vLLM from source:
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.11.0
-VLLM_USE_PRECOMPILED=1 uv pip install -e .
```
+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 .
+```
+
+### Development Environment for vLLM-Omni
Install vLLM-omni from source:
@@ -49,7 +65,7 @@ vLLM-omni's `pre-commit` hooks will now run automatically every time you commit.
```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
@@ -163,7 +179,7 @@ The goal of the vLLM-omni team is to be a _transparent reviewing machine_. We wo
## Additional Resources
-- [Design Documents](design_documents/index.md) - Architecture and design documentation
+- [Design Documents](../design/index.md) - Architecture and design documentation
## Thank You
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..ac9323cbc
--- /dev/null
+++ b/docs/contributing/model/README.md
@@ -0,0 +1 @@
+# Summary
diff --git a/docs/contributing/design_documents/api_design_doc.md b/docs/design/api_design_doc.md
similarity index 100%
rename from docs/contributing/design_documents/api_design_doc.md
rename to docs/design/api_design_doc.md
diff --git a/docs/contributing/design_documents/api_design_template.md b/docs/design/api_design_template.md
similarity index 100%
rename from docs/contributing/design_documents/api_design_template.md
rename to docs/design/api_design_template.md
diff --git a/docs/contributing/design_documents/architecture_overview.md b/docs/design/architecture_overview.md
similarity index 100%
rename from docs/contributing/design_documents/architecture_overview.md
rename to docs/design/architecture_overview.md
diff --git a/docs/contributing/design_documents/index.md b/docs/design/index.md
similarity index 100%
rename from docs/contributing/design_documents/index.md
rename to docs/design/index.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 100%
rename from docs/contributing/design_documents/vllm_omni_design.md
rename to docs/design/vllm_omni_design.md
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/features/README.md b/docs/features/README.md
new file mode 100644
index 000000000..c18fb056e
--- /dev/null
+++ b/docs/features/README.md
@@ -0,0 +1,3 @@
+# Features
+
+vLLM-Omni introduces new features to support multimodal comprehension and generation models inference and serving.
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..17f1b1f5c
--- /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#nvidia-cuda)
+- [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..5b7d080f7
--- /dev/null
+++ b/docs/getting_started/installation/gpu.md
@@ -0,0 +1,49 @@
+# GPU
+
+vLLM-Omni is a Python library that supports the following GPU variants. Select your GPU type to see vendor specific instructions:
+
+[](){ #nvidia-cuda }
+
+=== "NVIDIA CUDA"
+
+ --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:installation"
+
+## Requirements
+
+- OS: Linux
+- Python: 3.12
+
+!!! note
+ vLLM-Omni does not support Windows natively.
+
+=== "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"
+
+=== "NVIDIA CUDA"
+
+ --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:set-up-using-python"
+
+### 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
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..59087cee7
--- /dev/null
+++ b/docs/getting_started/installation/gpu/cuda.inc.md
@@ -0,0 +1,73 @@
+# --8<-- [start:installation]
+
+vLLM-Omni mainly contains python implementations for framework and models.
+
+# --8<-- [end:installation]
+# --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
+
+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-Omni
+
+```bash
+uv pip install vllm-omni
+```
+
+# --8<-- [end:pre-built-wheels]
+
+# --8<-- [start:build-wheel-from-source]
+
+## Installation of vLLM
+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 .
+```
+
+# --8<-- [end:build-wheel-from-source]
diff --git a/docs/getting_started/installation/npu.md b/docs/getting_started/installation/npu.md
new file mode 100644
index 000000000..d341c1504
--- /dev/null
+++ b/docs/getting_started/installation/npu.md
@@ -0,0 +1,3 @@
+# Ascend-NPU
+
+vLLM-Omni is a Python library that supports the following NPU variants. Select your NPU type to see vendor specific instructions:
diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md
new file mode 100644
index 000000000..c20460ed9
--- /dev/null
+++ b/docs/getting_started/installation/npu/npu.inc.md
@@ -0,0 +1,5 @@
+# --8<-- [start:installation]
+
+vLLM-Omni mainly contains python implementations for framework and models.
+
+# --8<-- [end:installation]
diff --git a/docs/getting_started/installation/python_env_setup.inc.md b/docs/getting_started/installation/python_env_setup.inc.md
new file mode 100644
index 000000000..06794f8d3
--- /dev/null
+++ b/docs/getting_started/installation/python_env_setup.inc.md
@@ -0,0 +1,6 @@
+It's recommended to use [uv](https://docs.astral.sh/uv/), a very fast Python environment manager, to create and manage Python environments. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) 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
+```
diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md
new file mode 100644
index 000000000..c8fdf7665
--- /dev/null
+++ b/docs/getting_started/quickstart.md
@@ -0,0 +1,23 @@
+# Quickstart
+
+This guide will help you quickly get started with vLLM-Omni to perform:
+
+- Offline batched inference
+- Online serving using OpenAI-compatible server
+
+## Prerequisites
+
+- OS: Linux
+- Python: 3.12
+
+## Installation
+
+Please refer to [installation](installation/README.md)
+
+## Offline Inference
+
+Please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
+
+## Online Serving with OpenAI-Completions API
+
+Please refer to [online serving](../user_guide/examples/online_serving/README.md)
diff --git a/docs/index.md b/docs/index.md
deleted file mode 100644
index 0db4f2908..000000000
--- a/docs/index.md
+++ /dev/null
@@ -1,32 +0,0 @@
-# Welcome to vLLM-omni
-
-**vLLM-omni** extends vLLM capabilities to support multi-modality model inference and serving, particularly focusing on non-autoregressive structures and non-textual outputs. It aims to support:
-
-- **Multi-modal Models**: Text, image, video, audio, and sensor data processing
-- **Non-autoregressive Architectures**: Diffusion Transformers (DiT) and other parallel generation models
-- **Heterogeneous Outputs**: Beyond traditional text generation to multimodal outputs
-
-## π― 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..f1819cfc8
--- /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/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)
+ -
+ - [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)`
+ - ``
+
+ 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.+?):(?P.+?)(#(?P.+?))?"
+ inline_link = re.compile(r"\[(?P[^\[]+?)\]\(" + 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..c2d91db5e
--- /dev/null
+++ b/docs/models/supported_models.md
@@ -0,0 +1,34 @@
+# 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>.
+
+## 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`|
+
+
+## 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..0e82bb9b0 100644
--- a/docs/user_guide/examples/offline_inference/qwen2_5_omni.md
+++ b/docs/user_guide/examples/offline_inference/qwen2_5_omni.md
@@ -1,6 +1,6 @@
# Offline Example of vLLM-omni for Qwen2.5-omni
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples\offline_inference\qwen2_5_omni>.
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni>.
## π οΈ Installation
@@ -36,29 +36,29 @@ Then run the command below.
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
+```
+
## Example materials
??? abstract "end2end.py"
``````py
- --8<-- "examples\offline_inference\qwen2_5_omni\end2end.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"
+ --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
``````
??? abstract "run_multiple_prompts.sh"
``````sh
- --8<-- "examples\offline_inference\qwen2_5_omni\run_multiple_prompts.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_single_prompt.sh"
``````
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
index 7b3341255..f6312fdae 100644
--- a/docs/user_guide/examples/online_serving/README.md
+++ b/docs/user_guide/examples/online_serving/README.md
@@ -1,6 +1,6 @@
# 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>.
+Source <https://github.com/vllm-project/vllm-omni/blob/main/examples/online_serving/README.md>.
## π οΈ Installation
@@ -26,10 +26,55 @@ cd examples/online_serving
Send request via python
```bash
-python openai_chat_completion_client_for_multimodal_generation.py
+python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
```
Send request via curl
```bash
-bash run_curl_multimodal_generation.sh
+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
+
+The gradio demo also supports running with an existing API server and can be customized with the following arguments.
+
+
+```bash
+python gradio_demo.py \
+ --model Qwen/Qwen2.5-Omni-7B \
+ --use-api-server \
+ --api-base http://localhost:8091/v1 \
+ --ip 127.0.0.1 \
+ --port 7861
+```
+
+- `--model`: Model name
+- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmniLLM locally.
+- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
+- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
+- `--port`: Port for Gradio server (default: 7861)
+- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
+- `--share`: Share the Gradio demo publicly (creates a public link)
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/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/mkdocs.yml b/mkdocs.yml
index b5a54c6b7..f5239388e 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,19 @@ 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 +79,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 +91,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 +141,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..513ee1317 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -56,12 +56,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",
]
diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py
index 67bcc5fad..d0f48432a 100644
--- a/vllm_omni/config/model.py
+++ b/vllm_omni/config/model.py
@@ -10,7 +10,28 @@ import vllm_omni.model_executor.models as me_models
@config
@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
class OmniModelConfig(ModelConfig):
- """Configuration for Omni models, extending the base ModelConfig."""
+ """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"
diff --git a/vllm_omni/core/sched/diffusion_scheduler.py b/vllm_omni/core/sched/diffusion_scheduler.py
index 0474f83ab..1b960f937 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
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/engine/__init__.py b/vllm_omni/engine/__init__.py
index 2de831388..6a7f48525 100644
--- a/vllm_omni/engine/__init__.py
+++ b/vllm_omni/engine/__init__.py
@@ -49,6 +49,19 @@ 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)
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 64df79c7b..6e7b02be4 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -9,12 +9,33 @@ from vllm_omni.config import OmniModelConfig
@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
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()
@@ -37,12 +58,33 @@ 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
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()
diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py
index fb664b7a2..97bab2633 100644
--- a/vllm_omni/engine/output_processor.py
+++ b/vllm_omni/engine/output_processor.py
@@ -19,6 +19,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,
@@ -86,6 +93,17 @@ class OmniRequestState(RequestState):
)
def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], mm_type: Optional[str]) -> None:
+ """Add a multimodal tensor to the accumulated outputs.
+
+ Accumulates multimodal tensors (e.g., image latents, audio waveforms)
+ that are produced incrementally during generation. Tensors are
+ concatenated along the first dimension.
+
+ Args:
+ tensor: Optional tensor to add. If None, this is a no-op.
+ mm_type: Optional modality type identifier (e.g., "image", "audio").
+ Used to categorize the output type.
+ """
if tensor is None:
return
try:
@@ -115,6 +133,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
@@ -179,6 +214,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 +230,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 +250,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 +289,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()
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..7cc138efb 100644
--- a/vllm_omni/entrypoints/async_omni_llm.py
+++ b/vllm_omni/entrypoints/async_omni_llm.py
@@ -53,6 +53,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,
@@ -63,7 +95,7 @@ class AsyncOmniLLM(EngineClient):
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)
@@ -140,6 +172,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)
@@ -162,7 +200,11 @@ class AsyncOmniLLM(EngineClient):
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 +220,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
@@ -510,9 +579,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 +635,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/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py
index 655d0307f..c03d3ce2f 100644
--- a/vllm_omni/entrypoints/omni_llm.py
+++ b/vllm_omni/entrypoints/omni_llm.py
@@ -42,11 +42,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 +84,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)
@@ -132,6 +157,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 +188,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:
@@ -465,13 +517,31 @@ class OmniLLM:
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,
+ **kwargs: Any,
):
"""LLM constructor."""
if "disable_log_stats" not in kwargs:
diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py
index ad89b4cbb..c52f77043 100644
--- a/vllm_omni/entrypoints/omni_stage.py
+++ b/vllm_omni/entrypoints/omni_stage.py
@@ -39,7 +39,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 +87,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 +163,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
@@ -156,6 +217,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)
@@ -173,10 +239,22 @@ class OmniStage:
self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
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()
@@ -184,9 +262,24 @@ 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:
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/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/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
index 54d2b2359..e4f6b8e7b 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
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/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/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_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_model_runner.py b/vllm_omni/worker/gpu_model_runner.py
index 55e5e1391..e9de36726 100644
--- a/vllm_omni/worker/gpu_model_runner.py
+++ b/vllm_omni/worker/gpu_model_runner.py
@@ -30,7 +30,26 @@ logger = init_logger(__name__)
class OmniGPUModelRunner(GPUModelRunner):
+ """GPU model runner for omni models with M-RoPE support.
+
+ Extends the base GPUModelRunner to handle multimodal rotary position
+ embeddings (M-RoPE) for processing image, video, and audio inputs
+ with proper positional encoding.
+ """
+
def _init_mrope_positions(self, req_state: CachedRequestState):
+ """Initialize M-RoPE positions for multimodal inputs.
+
+ Extracts multimodal feature metadata (image grids, video grids,
+ audio features) and computes M-RoPE positions for proper positional
+ encoding of multimodal tokens.
+
+ Args:
+ req_state: Cached request state containing multimodal features
+
+ Raises:
+ AssertionError: If the model does not support M-RoPE
+ """
image_grid_thw = []
video_grid_thw = []
second_per_grid_ts = []
+ 
+