update

.buildkite/ci_config.yaml

@@ -1,10 +0,0 @@
-name: vllm_omni_ci
-github_repo_name: vllm-project/vllm-omni
-job_dirs:
-  - ".buildkite/jobs"
-run_all_patterns: []
-run_all_exclude_patterns: []
-registries: public.ecr.aws/q9t5s3a7
-repositories:
-  main: "vllm-ci-postmerge-repo"
-  premerge: "vllm-ci-test-repo"

.buildkite/jobs/build.yaml

@@ -1,10 +0,0 @@
-group: Build
-steps:
-  - label: ":docker: build vllm-omni-ci image"
-    key: image-build
-    depends_on: []
-    commands:
-      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
-      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"

.buildkite/jobs/tests.yaml

@@ -1,30 +0,0 @@
-group: Tests
-depends_on:
-  - image-build
-steps:
-- label: "Simple Unit Test"
-  commands:
-    - ".buildkite/scripts/simple_test.sh"
-  no_gpu: true
-  no_plugin: true
-
-- label: "Diffusion Model Test"
-  timeout_in_minutes: 15
-  commands:
-    - pytest -s -v tests/single_stage/test_diffusion_model.py
-
-- label: "Omni Model Test"
-  timeout_in_minutes: 15
-  num_gpus: 4
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/multi_stages/
-
-- label: "Omni Model Test with H100"
-  timeout_in_minutes: 20
-  gpu: h100
-  num_gpus: 2
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/multi_stages_h100/

.buildkite/pipeline.yml

@@ -0,0 +1,6 @@
+steps:
+  - label: "Simple Unit Test"
+    commands:
+      - ".buildkite/scripts/simple_test.sh"
+    agents:
+      queue: "cpu_queue_premerge"

.github/ISSUE_TEMPLATE/config.yml

@@ -1,5 +1,5 @@
 blank_issues_enabled: false
 contact_links:
   - name: Questions
-    url: https://slack.vllm.ai
-    about: Ask questions and discuss with other vLLM-Omni community members in our #sig-omni slack channel.
+    url: https://github.com/vllm-project/vllm-omni/discussions
+    about: Ask questions and discuss with other vLLM-Omni community members

.gitignore

@@ -156,9 +156,15 @@ cython_debug/
 # cursor
 .cursor/
 
+# docker
+docker/
+
 # scripts
 /scripts/
 
+# tests
+tests/
+
 # PyCharm
 #  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
 #  be added to the global gitignore or merged into this project gitignore.  For a PyCharm

CONTRIBUTING.md

@@ -1,3 +1,3 @@
 # Contributing to vLLM-Omni
 
-You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/)
+You may find information about contributing to vLLM-Omni on [Contributing](docs/contributing/)

README.md

@@ -1,7 +1,7 @@
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
-    <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="./docs/source/logos/vllm-omni-logo-text-dark.png">
+    <img alt="vllm-omni" src="./docs/source/logos/vllm-omni-logo-text-dark.png" width=55%>
   </picture>
 </p>
 <h3 align="center">
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://tinyurl.com/vllm-omni-meeting"><b>Weekly Meeting</b></a> |
 </p>
 
 ---
@@ -30,7 +30,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 <p align="center">
   <picture>
-    <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/omni-modality-model-architecture.png" width=55%>
+    <img alt="vllm-omni" src="./docs/source/architecture/omni-modality model architecture.png" width=55%>
   </picture>
 </p>
 
@@ -57,6 +57,8 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 
 Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn more.
 
+Visit our documentation to learn more.
+
 - [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/)
@@ -66,8 +68,10 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
 We welcome and value any contributions and collaborations.
 Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.
 
-## Join the Community
-Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).
+## Weekly Meeting
+
+- vLLM Omni Weekly Meeting: https://tinyurl.com/vllm-omni-meeting
+- Wednesday, 11:30 - 12:30 (UTC+8, [Convert to your timezone](https://dateful.com/convert/gmt8?t=15))
 
 ## License
 

docker/Dockerfile.ci

@@ -1,14 +0,0 @@
-ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.11.0
-FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
-ARG APP_DIR=/workspace/vllm-omni
-WORKDIR ${APP_DIR}
-
-COPY . .
-
-# Install vllm-omni into the same uv-managed Python environment used by the base image.
-RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
-
-RUN ln -sf /usr/bin/python3 /usr/bin/python
-
-ENTRYPOINT []

docs/.nav.yml

@@ -1,51 +1,44 @@
 nav:
-- Home: README.md
-- User Guide:
-  - Getting Started:
-    - getting_started/quickstart.md
-    - getting_started/installation
-  - Examples:
-    - examples/README.md
-    - Offline Inference:
-      - Offline Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
-      - Offline Example of vLLM-Omni for Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
-      - Qwen-Image Offline Inference: user_guide/examples/offline_inference/qwen_image.md
-    - Online Serving:
-      - Online serving Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
-      - Online serving Example of vLLM-Omni for Qwen3-omni: user_guide/examples/online_serving/qwen3_omni.md
-  - General:
-    - usage/*
-  - Configuration:
-    - configuration/README.md
-    - configuration/*
-  - Models:
-    - models/supported_models.md
-- Developer Guide:
-  - General:
-    - contributing/README.md
-    - glob: contributing/*
-      flatten_single_child_sections: true
-  - Model Implementation:
-    - contributing/model/README.md
-  - CI: contributing/ci
-  - Design Documents:
-    - design/index.md
-    - design/architecture_overview.md
-    - design/vllm_omni_design.md
-    - Feature Design:
-      - design/feature/multi_request_streaming_design.md
-      - design/feature/omni_connector_design.md
-    - Module Design:
-      - design/module/ar_module_design.md
-      - design/module/dit_module_design.md
-      - design/module/entrypoint_module_design.md
-  - Docs Guide: contributing/DOCS_GUIDE.md
-- API Reference:
-  - api/README.md
-  - api/vllm_omni
-- CLI Reference: cli
-- Community:
-  - community/*
-  - Slack: https://slack.vllm.ai
-  - Blog: https://blog.vllm.ai
-  - Forum: https://discuss.vllm.ai
+  - Home: README.md
+  - User Guide:
+    - Getting Started:
+      - getting_started/quickstart.md
+      - getting_started/installation
+    - Examples:
+      - examples/README.md
+      - Offline Inference:
+        - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
+      - Online Serving:
+        - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
+        - Qwen2.5-Image: user_guide/examples/offline_inference/qwen_image.md
+    - General:
+      - usage/*
+    - Configuration:
+      - configuration/README.md
+      - configuration/*
+    - Models:
+      - models/supported_models.md
+  - Developer Guide:
+    - General:
+      - contributing/README.md
+      - glob: contributing/*
+        flatten_single_child_sections: true
+    - Model Implementation:
+      - contributing/model/README.md
+    - CI: contributing/ci
+    - Design Documents:
+      - design/index.md
+      - design/architecture_overview.md
+      - design/vllm_omni_design.md
+      - design/mrs_design.md
+      - design/api_design_doc.md
+    - Docs Guide: contributing/DOCS_GUIDE.md
+  - API Reference:
+    - api/README.md
+    - api/vllm_omni
+  - CLI Reference: cli
+  - Community:
+    - community/*
+    - Slack: https://slack.vllm.ai
+    - Blog: https://blog.vllm.ai
+    - Forum: https://discuss.vllm.ai

docs/README.md

@@ -8,8 +8,8 @@ hide:
 
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="./source/logos/vllm-omni-logo.png">
-    <img alt="vllm-omni" src="./source/logos/vllm-omni-logo.png" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="./source/logos/vllm-omni-logo-text-dark.png">
+    <img alt="vllm-omni" src="./source/logos/vllm-omni-logo-text-dark.png" width=55%>
   </picture>
 </p>
 <h3 align="center">
@@ -34,8 +34,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="./source/architecture/omni-modality-model-architecture.png">
-    <img alt="vllm-omni-arch" src="./source/architecture/omni-modality-model-architecture.png" width=55%>
+    <img alt="vllm-omni" src="./source/architecture/omni-modality model architecture.png" width=55%>
   </picture>
 </p>
 
@@ -59,6 +58,4 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 - Multi-modality generation models (e.g. Qwen-Image)
 
 For more information, checkout the following:
-
-- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1Y7t2Zm3BIISPN-_X_sDpxZMNCWir2rAx/edit?usp=drive_link&ouid=110473603432222024453&rtpof=true&sd=true)
-- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
+- [Placeholder for vllm-omni blogpost]()

docs/api/README.md

@@ -4,8 +4,8 @@
 
 Main entry points for vLLM-Omni inference and serving.
 
-- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
-- [vllm_omni.entrypoints.async_omni.AsyncOmniStageLLM][]
+- [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][]
@@ -34,8 +34,6 @@ Engine classes for offline and online inference.
 - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
-- [vllm_omni.engine.OmniEngineCoreOutput][]
-- [vllm_omni.engine.OmniEngineCoreOutputs][]
 - [vllm_omni.engine.OmniEngineCoreRequest][]
 - [vllm_omni.engine.PromptEmbedsPayload][]
 - [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
@@ -49,15 +47,15 @@ Engine classes for offline and online inference.
 Core scheduling and caching components.
 
 - [vllm_omni.core.dit_cache_manager.DiTCacheManager][]
-- [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][]
-- [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
+- [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.output_templates.OmniOutput][]
+- [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][]
@@ -74,22 +72,6 @@ Model execution components.
 - [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][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_moe.Qwen3MoeForCausalLM][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni.Qwen3OmniMoeForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav.Qwen3OmniMoeCode2Wav][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniCodePredictorBaseModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniMoeTalkerCodePredictor][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMForCausalLM][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeConditionalGenerationMixin][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerMultiModalProcessor][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerProcessingInfo][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3Omni_VisionTransformer][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3_VisionPatchEmbed][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3_VisionPatchMerger][]
 
 ## Configuration
 
@@ -103,15 +85,8 @@ Worker classes and model runners for distributed inference.
 
 - [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][]
 - [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][]
-- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][]
-- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
-- [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
-- [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
+- [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
+- [vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
 - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
-- [vllm_omni.worker.npu.npu_ar_model_runner.NPUARModelRunner][]
-- [vllm_omni.worker.npu.npu_ar_worker.NPUARWorker][]
-- [vllm_omni.worker.npu.npu_generation_model_runner.NPUGenerationModelRunner][]
-- [vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker][]
-- [vllm_omni.worker.npu.npu_model_runner.OmniNPUModelRunner][]

docs/community/contact_us.md

@@ -2,4 +2,4 @@
 
 - For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
 - For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/)
-- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm-omni/security/advisories) feature
+- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm/security/advisories) feature

docs/community/volunteers.md

@@ -1,12 +0,0 @@
-# Volunteers for Bugfix and CI
-
-We encourage you to check current docs and [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions for your questions. If non of these can solve it, please propose an issue to describe your questions about bug or CI problems for developing.
-
-If you have urgent need for locating and solving bugfix or CI problems, please find community volunteers below.
-
-| Dec 4-Dec 12 | Dec 15-Dec 19 | Dec 22-Dec 26 | Dec 29- Jan 2, 2026| Jan 5-Jan 9 | Jan 12-Jan 16 |
-|----------|----------|----------|----------|----------|----------|
-| <a href="https://github.com/congw729">Conw729</a> | <a href="https://github.com/yinpeiqi">yinpeiqi</a> | <a href="https://github.com/tzhouam">tzhouam</a> | <a href="https://github.com/SamitHuang">SamitHuang</a> | <a href="https://github.com/gcanlin">gcanlin</a> | <a href="https://github.com/natureofnature">natureofnature</a> |
-| <a href="https://github.com/david6666666">david6666666</a> | <a href="https://github.com/R2-Y">R2-Y</a> | <a href="https://github.com/hsliuustc0106">hsliuustc0106</a> | <a href="https://github.com/Gaohan123">Gaohan123</a> | <a href="https://github.com/ZJY0516">ZJY0516</a> | <a href="https://github.com/qibaoyuan">qibaoyuan</a> |
-
-We kindly welcome more contributors to fix bugs and contribute new features!

docs/configuration/stage_configs.md

@@ -2,9 +2,6 @@
 
 In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion.
 
-!!! note
-    Default stage config YAMLs (for example, `vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml` and `vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml`) are bundled and loaded automatically when `stage_configs_path` is not provided. They have been verified to work on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni.
-
 Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions:
 
 - Claim partition of stages and their corresponding class implementation in `model_executor/models`.
@@ -47,7 +44,7 @@ stage_args:
       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 AsyncOmni will use its tokenizer as default
+    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.
@@ -209,7 +206,7 @@ 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, `AsyncOmni` will use its tokenizer as the default tokenizer for the entire pipeline.
+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`
 

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -9,7 +9,7 @@ stage_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      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
@@ -35,7 +35,7 @@ stage_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
       trust_remote_code: true

docs/contributing/DOCS_GUIDE.md

@@ -9,7 +9,7 @@ This directory contains the source files for the vLLM-Omni documentation.
 Install documentation dependencies:
 
 ```bash
-uv pip install -e ".[docs]"
+pip install -e ".[docs]"
 ```
 
 ### Build and Serve Documentation
@@ -84,7 +84,7 @@ The documentation is automatically deployed to GitHub Pages using GitHub Actions
    ```
 
 3. **Documentation will be available at**:
-   - `https://vllm-omni.readthedocs.io`
+   - `https://vllm-project.github.io/vllm-omni/`
 
 The GitHub Actions workflow (`.github/workflows/docs.yml`) will automatically:
 - Build the documentation when you push to `main` branch

docs/contributing/README.md

@@ -2,9 +2,6 @@
 
 Thank you for your interest in contributing to vLLM-Omni! This document provides guidelines and instructions for contributing.
 
-!!! note
-    We host weekly developer-facing online meetings to discuss milestones and updates **every Tuesday at 19:30 PDT**. Meeting link as well as the past meeting notes can be found [here](https://tinyurl.com/vllm-omni-meeting).
-
 ## Getting Started
 
 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:
@@ -23,7 +20,7 @@ Install vLLM-Omni from source with development dependencies:
 
 ```bash
 git clone https://github.com/vllm-project/vllm-omni.git
-cd vllm-omni
+cd vllm_omni
 uv pip install -e ".[dev]"
 ```
 

docs/design/api_design_doc.md

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

docs/design/api_design_template.md

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

docs/design/feature/omni_connector_design.md

@@ -1 +0,0 @@
-## OmniConnector for fully disaggregated serving (update soon)

docs/design/index.md

@@ -6,14 +6,8 @@ This section contains design documents and architecture specifications for vLLM-
 
 - [Architecture Overview](architecture_overview.md)
 - [vLLM-Omni Software Design Document](vllm_omni_design.md)
+- [Multi-Request Streaming](mrs_design.md)
 
-## Feature Design Documents
+## API Design Documentation
 
-- [Multi-Request Streaming](feature/multi_request_streaming_design.md)
-- [Omni Connector Design](feature/omni_connector_design.md)
-
-## Module Design Documents
-
-- [AR Module Design](module/ar_module_design.md)
-- [DIT Module Design](module/dit_module_design.md)
-- [Entrypoint Module Design](module/entrypoint_module_design.md)
+- [vLLM-Omni API Design Documentation](api_design_doc.md)

docs/design/module/ar_module_design.md

@@ -1 +0,0 @@
-Architecture design of the AR module (update soon)

docs/design/module/dit_module_design.md

@@ -1 +0,0 @@
-Architecture design of the DiT module (update soon)

docs/design/module/entrypoint_module_design.md

@@ -1 +0,0 @@
-Architecture design of the entrypoint (update soon)

docs/design/vllm_omni_design.md

@@ -14,9 +14,9 @@ vLLM-Omni is a multi-modality extension for vLLM that supports non-autoregressiv
 
 ## Key Data Flow
 
-API Server --> OmniLLM/AsyncOmni (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
+API Server --> OmniLLM/AsyncOmniLLM (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
  --> Scheduler (New one for DiT) --> Executor (New one for diffusers) --> Worker (New one for DiT)
- --> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessor (New one for final multimodal output)
+ --> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessoer (New one for final multimodal output)
 
 ## Core Components (Ordered by Data Flow)
 
@@ -40,7 +40,7 @@ graph TD
     B --> C{Detect --omni flag}
     C -->|Yes| D[Parse OmniConfig]
     C -->|No| E[Forward to vLLM CLI]
-    D --> F[Initialize AsyncOmni]
+    D --> F[Initialize AsyncOmniLLM]
     F --> G[Start omni Server]
     G --> H[Multi-stage Processing]
     E --> I[Standard vLLM Pipeline]
@@ -137,7 +137,7 @@ class OmniLLM(LLM):
             prompt_str, engine_request, tokenization_kwargs = self._process_stage_inputs(stage_config, **stage_args)
 
             # Add inputs to Engine
-            stage_engine.add_request(request_id, prompt_str, tokenization_kwargs)
+            stage_engine.add_request(requesy_id, prompt_str, tokenization_kwargs)
             # Run Engine
             stage_output = stage_engine.step()
 
@@ -157,7 +157,7 @@ class OmniLLM(LLM):
             return self._process_ar_inputs(**stage_args)
         elif stage_config.engine_type == "DiT":
             return self._process_dit_inputs(**stage_args)
-        else:
+        else
             raise NotImplementedError
 
     def _process_dit_inputs(**stage_args)-> Any:
@@ -209,7 +209,7 @@ Similar to OmniLLM in offline inference, add some asynchronous processing, refer
 from vllm.v1.engine.async_llm import AsyncLLM
 
 
-class AsyncOmni(AsyncLLM):
+class AsyncOmniLLM(AsyncLLM):
     """Extended AsyncLLM supporting multiple engines and stage-based processing"""
 
     def __init__(self, stage_configs: List[StageConfig]):
@@ -449,15 +449,15 @@ AsyncLLM.add_request()
 ↓
 OutputProcessor.add_request()
 ↓
-RequestState.from_new_request() → create request state
+RequestState.from_new_request() → 创建请求状态
 
-AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() →  create a background loop that continuously pulls output from EngineCore
+AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() →  创建一个Background loop 持续从EngineCore获取输出
 ↓
-OutputProcessor.process_outputs() → update state and process outputs
+OutputProcessor.process_outputs() → 更新状态并处理输出
 ↓
-RequestState.make_request_output() → convert to the final output, formatted as RequestOutput or PoolingRequestOutput
+RequestState.make_request_output() → 转换为最终输出,格式为RequestOutput或者PoolingRequestOutput
 ↓
-RequestOutputCollector.put() → push to the queue (AsyncLLM)
+RequestOutputCollector.put() → 推送到队列（AsyncLLM）
 ```
 
 Need to add implementation for an existing method

docs/getting_started/installation/README.md

@@ -1,7 +1,7 @@
 # Installation
 
-vLLM-Omni supports the following hardware platforms:
+vLLM supports the following hardware platforms:
 
 - [GPU](gpu.md)
     - [NVIDIA CUDA](gpu.md)
-- [NPU](npu.md)
+- [Ascend NPU](npu.md)

docs/getting_started/installation/gpu.md

@@ -37,7 +37,4 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 ## Set up using Docker
 
 ### Pre-built images
-
-=== "NVIDIA CUDA"
-
-    --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-images"
+To be released... Stay tuned!

docs/getting_started/installation/gpu/cuda.inc.md

@@ -44,7 +44,7 @@ uv pip install vllm==0.11.0 --torch-backend=auto
 Install additional requirements for vLLM-Omni
 ```bash
 git clone https://github.com/vllm-project/vllm-omni.git
-cd vllm-omni
+cd vllm_omni
 uv pip install -e .
 ```
 
@@ -73,24 +73,3 @@ uv pip install --no-build-isolation --editable .
 </details>
 
 # --8<-- [end:build-wheel-from-source]
-
-
-# --8<-- [start:pre-built-images]
-
-vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which release of vLLM it is based on.
-
-Here's an example deployment command that has been verified on 2 x H100's:
-```bash
-docker run --runtime nvidia --gpus 2 \
-    -v ~/.cache/huggingface:/root/.cache/huggingface \
-    --env "HF_TOKEN=$HF_TOKEN" \
-    -p 8091:8091 \
-    --ipc=host \
-    vllm/vllm-omni:v0.11.0rc1 \
-    --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
-```
-
-!!! tip
-    You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
-
-# --8<-- [end:pre-built-images]

docs/getting_started/installation/npu.md

@@ -1,23 +1,3 @@
-# NPU
+# Ascend-NPU
 
-vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a community maintained hardware plugin for running vLLM on NPU.
-
-## Requirements
-
-- OS: Linux
-- Python: 3.12
-
-!!! note
-    vLLM-Omni is currently not natively supported on Windows.
-
-=== "NPU"
-
-    --8<-- "docs/getting_started/installation/npu/npu.inc.md:requirements"
-
-## Installation
-
-### Recommended
-
-=== "NPU"
-
-    --8<-- "docs/getting_started/installation/npu/npu.inc.md:installation"
+vLLM-Omni is a Python library that supports the following NPU variants. Select your NPU type to see vendor specific instructions:

docs/getting_started/installation/npu/npu.inc.md

@@ -1,49 +1,5 @@
-# --8<-- [start:requirements]
-
-For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
-
-# --8<-- [end:requirements]
 # --8<-- [start:installation]
 
-The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images:
-
-```bash
-# Update DEVICE according to your NPUs (/dev/davinci[0-7])
-export DEVICE0=/dev/davinci0
-export DEVICE1=/dev/davinci1
-# Update the vllm-ascend image
-# Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
-# Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
-docker run --rm \
-    --name vllm-omni-npu \
-    --shm-size=1g \
-    --device $DEVICE0 \
-    --device $DEVICE1 \
-    --device /dev/davinci_manager \
-    --device /dev/devmm_svm \
-    --device /dev/hisi_hdc \
-    -v /usr/local/dcmi:/usr/local/dcmi \
-    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
-    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-    -v /etc/ascend_install.info:/etc/ascend_install.info \
-    -v /root/.cache:/root/.cache \
-    -p 8000:8000 \
-    -it $IMAGE bash
-
-# Inside the container, install vLLM-Omni from source
-cd /vllm-workspace
-git clone https://github.com/vllm-project/vllm-omni.git
-cd vllm-omni
-pip install -v -e .
-export VLLM_WORKER_MULTIPROC_METHOD=spawn
-```
-
-The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code placed in `/vllm-workspace` installed in development mode.
-
-For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
+vLLM-Omni mainly contains python implementations for framework and models.
 
 # --8<-- [end:installation]

docs/getting_started/quickstart.md

@@ -16,19 +16,7 @@ Please refer to [installation](installation/README.md)
 
 ## Offline Inference
 
-Text-to-image generation quickstart with vLLM-Omni:
-
-```python
-from vllm_omni.entrypoints.omni import Omni
-
-if __name__ == "__main__":
-    omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
-    prompt = "a cup of coffee on the table"
-    images = omni.generate(prompt)
-    images[0].save("coffee.png")
-```
-
-For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
+Please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
 
 ## Online Serving with OpenAI-Completions API
 

docs/mkdocs/hooks/generate_examples.py

@@ -1,336 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import itertools
-import logging
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Literal
-
-import regex as re
-import yaml
-
-logger = logging.getLogger("mkdocs")
-
-ROOT_DIR = Path(__file__).parent.parent.parent.parent
-ROOT_DIR_RELATIVE = "../../../../.."
-EXAMPLE_DIR = ROOT_DIR / "examples"
-EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
-NAV_FILE = ROOT_DIR / "docs/.nav.yml"
-
-
-def fix_case(text: str) -> str:
-    subs = {
-        "api": "API",
-        "cli": "CLI",
-        "cpu": "CPU",
-        "llm": "LLM",
-        "mae": "MAE",
-        "tpu": "TPU",
-        "gguf": "GGUF",
-        "lora": "LoRA",
-        "rlhf": "RLHF",
-        "vllm": "vLLM",
-        "openai": "OpenAI",
-        "lmcache": "LMCache",
-        "multilora": "MultiLoRA",
-        "mlpspeculator": "MLPSpeculator",
-        r"fp\d+": lambda x: x.group(0).upper(),  # e.g. fp16, fp32
-        r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
-    }
-    for pattern, repl in subs.items():
-        text = re.sub(rf"\b{pattern}\b", repl, text, flags=re.IGNORECASE)
-    return text
-
-
-@dataclass
-class Example:
-    """
-    Example class for generating documentation content from a given path.
-
-    Attributes:
-        path (Path): The path to the main directory or file.
-        category (str): The category of the document.
-        main_file (Path): The main file in the directory.
-        other_files (list[Path]): list of other files in the directory.
-        title (str): The title of the document.
-
-    Methods:
-        __post_init__(): Initializes the main_file, other_files, and title attributes.
-        determine_main_file() -> Path: Determines the main file in the given path.
-        determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
-        determine_title() -> str: Determines the title of the document.
-        generate() -> str: Generates the documentation content.
-    """  # noqa: E501
-
-    path: Path
-    category: str = None
-    main_file: Path = field(init=False)
-    other_files: list[Path] = field(init=False)
-    title: str = field(init=False)
-
-    def __post_init__(self):
-        self.main_file = self.determine_main_file()
-        self.other_files = self.determine_other_files()
-        self.title = self.determine_title()
-
-    @property
-    def is_code(self) -> bool:
-        return self.main_file.suffix != ".md"
-
-    def determine_main_file(self) -> Path:
-        """
-        Determines the main file in the given path.
-        If the path is a file, it returns the path itself. Otherwise, it searches
-        for Markdown files (*.md) in the directory and returns the first one found.
-        Returns:
-            Path: The main file path, either the original path if it's a file or the first
-            Markdown file found in the directory.
-        Raises:
-            IndexError: If no Markdown files are found in the directory.
-        """  # noqa: E501
-        return self.path if self.path.is_file() else list(self.path.glob("*.md")).pop()
-
-    def determine_other_files(self) -> list[Path]:
-        """
-        Determine other files in the directory excluding the main file.
-
-        This method checks if the given path is a file. If it is, it returns an empty list.
-        Otherwise, it recursively searches through the directory and returns a list of all
-        files that are not the main file.
-
-        Returns:
-            list[Path]: A list of Path objects representing the other files in the directory.
-        """  # noqa: E501
-        if self.path.is_file():
-            return []
-        # Binary file extensions to exclude
-        binary_extensions = {
-            ".wav",
-            ".mp3",
-            ".mp4",
-            ".avi",
-            ".mov",
-            ".mkv",  # Audio/Video
-            ".png",
-            ".jpg",
-            ".jpeg",
-            ".gif",
-            ".bmp",
-            ".ico",
-            ".svg",  # Images
-            ".pdf",
-            ".zip",
-            ".tar",
-            ".gz",
-            ".bz2",
-            ".xz",  # Archives/Documents
-            ".exe",
-            ".so",
-            ".dll",
-            ".dylib",  # Binaries
-            ".bin",
-            ".dat",
-            ".db",
-            ".sqlite",  # Data files
-        }
-
-        def is_other_file(file: Path) -> bool:
-            return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions
-
-        return [file for file in self.path.rglob("*") if is_other_file(file)]
-
-    def determine_title(self) -> str:
-        if not self.is_code:
-            # Specify encoding for building on Windows
-            with open(self.main_file, encoding="utf-8") as f:
-                first_line = f.readline().strip()
-            match = re.match(r"^#\s+(?P<title>.+)$", first_line)
-            if match:
-                return match.group("title")
-        return fix_case(self.path.stem.replace("_", " ").title())
-
-    def fix_relative_links(self, content: str) -> str:
-        """
-        Fix relative links in markdown content by converting them to gh-file
-        format.
-
-        Args:
-            content (str): The markdown content to process
-
-        Returns:
-            str: Content with relative links converted to gh-file format
-        """
-        # Regex to match markdown links [text](relative_path)
-        # This matches links that don't start with http, https, ftp, or #
-        link_pattern = r"\[([^\]]*)\]\((?!(?:https?|ftp)://|#)([^)]+)\)"
-
-        def replace_link(match):
-            link_text = match.group(1)
-            relative_path = match.group(2)
-
-            # Make relative to repo root
-            gh_file = (self.main_file.parent / relative_path).resolve()
-            gh_file = gh_file.relative_to(ROOT_DIR)
-
-            # Make GitHub URL
-            url = "https://github.com/vllm-project/vllm-omni/"
-            url += "tree/main" if self.path.is_dir() else "blob/main"
-            gh_url = f"{url}/{gh_file}"
-
-            return f"[{link_text}]({gh_url})"
-
-        return re.sub(link_pattern, replace_link, content)
-
-    def generate(self) -> str:
-        content = f"# {self.title}\n\n"
-        url = "https://github.com/vllm-project/vllm-omni/"
-        url += "tree/main" if self.path.is_dir() else "blob/main"
-        content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
-
-        # Use long code fence to avoid issues with
-        # included files containing code fences too
-        code_fence = "``````"
-
-        if self.is_code:
-            main_file_rel = self.main_file.relative_to(ROOT_DIR)
-            content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
-        else:
-            with open(self.main_file) as f:
-                # Skip the title from md snippets as it's been included above
-                main_content = f.readlines()[1:]
-            content += self.fix_relative_links("".join(main_content))
-        content += "\n"
-
-        if not self.other_files:
-            return content
-
-        content += "## Example materials\n\n"
-        for file in sorted(self.other_files):
-            content += f'??? abstract "{file.relative_to(self.path)}"\n'
-            if file.suffix != ".md":
-                content += f"    {code_fence}{file.suffix[1:]}\n"
-            content += f'    --8<-- "{file.relative_to(ROOT_DIR)}"\n'
-            if file.suffix != ".md":
-                content += f"    {code_fence}\n"
-
-        return content
-
-
-def update_nav_file(examples: list[Example]):
-    """
-    Update the .nav.yml file to include all generated examples.
-    This function completely regenerates the examples section based on the actual
-    folder structure, ensuring consistency between the examples folder and nav file.
-
-    Args:
-        examples: List of Example objects that have been generated
-    """
-    if not NAV_FILE.exists():
-        logger.warning("Navigation file not found: %s", NAV_FILE)
-        return
-
-    # Read the current nav file
-    with open(NAV_FILE, encoding="utf-8") as f:
-        nav_data = yaml.safe_load(f) or {}
-
-    nav_list = nav_data.get("nav", [])
-
-    # Find the "User Guide" section
-    user_guide_idx = None
-    examples_idx = None
-    for i, item in enumerate(nav_list):
-        if isinstance(item, dict) and "User Guide" in item:
-            user_guide_idx = i
-            user_guide_content = item["User Guide"]
-            # Find the "Examples" subsection
-            for j, subitem in enumerate(user_guide_content):
-                if isinstance(subitem, dict) and "Examples" in subitem:
-                    examples_idx = j
-                    break
-            break
-
-    if user_guide_idx is None or examples_idx is None:
-        logger.warning("Could not find 'User Guide' -> 'Examples' section in nav file")
-        return
-
-    # Get existing Examples section to preserve non-example items (like README.md)
-    existing_examples_content = nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"]
-
-    # Preserve string items (like "examples/README.md") that are not example categories
-    preserved_items = [
-        item
-        for item in existing_examples_content
-        if isinstance(item, str) and not item.startswith("user_guide/examples/")
-    ]
-
-    # Group examples by category
-    examples_by_category = {}
-    for example in examples:
-        category = example.category
-        if category not in examples_by_category:
-            examples_by_category[category] = []
-        examples_by_category[category].append(example)
-
-    # Build the new Examples section - start with preserved items
-    examples_section = preserved_items.copy()
-
-    # Add examples grouped by category, sorted by category name
-    for category in sorted(examples_by_category.keys()):
-        category_examples = sorted(examples_by_category[category], key=lambda e: e.path.stem)
-        category_items = []
-        for example in category_examples:
-            doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
-            rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            category_items.append({example.title: str(rel_path)})
-
-        if category_items:
-            # Format category name (e.g., "offline_inference" -> "Offline Inference")
-            category_title = fix_case(category.replace("_", " ").title())
-            examples_section.append({category_title: category_items})
-
-    # Update the nav structure
-    nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section
-
-    # Write back to file
-    nav_data["nav"] = nav_list
-    with open(NAV_FILE, "w", encoding="utf-8") as f:
-        yaml.dump(nav_data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
-    logger.info("Updated navigation file: %s", NAV_FILE.relative_to(ROOT_DIR))
-
-
-def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
-    logger.info("Generating example documentation")
-    logger.debug("Root directory: %s", ROOT_DIR.resolve())
-    logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
-    logger.debug("Example document directory: %s", EXAMPLE_DOC_DIR.resolve())
-
-    # Create the EXAMPLE_DOC_DIR if it doesn't exist
-    if not EXAMPLE_DOC_DIR.exists():
-        EXAMPLE_DOC_DIR.mkdir(parents=True)
-
-    categories = sorted(p for p in EXAMPLE_DIR.iterdir() if p.is_dir())
-
-    examples = []
-    glob_patterns = ["*.py", "*.md", "*.sh"]
-    # Find categorised examples
-    for category in categories:
-        globs = [category.glob(pattern) for pattern in glob_patterns]
-        for path in itertools.chain(*globs):
-            examples.append(Example(path, category.stem))
-        # Find examples in subdirectories
-        for path in category.glob("*/*.md"):
-            examples.append(Example(path.parent, category.stem))
-
-    # Generate the example documentation
-    for example in sorted(examples, key=lambda e: e.path.stem):
-        example_name = f"{example.path.stem}.md"
-        doc_path = EXAMPLE_DOC_DIR / example.category / example_name
-        if not doc_path.parent.exists():
-            doc_path.parent.mkdir(parents=True)
-        # Specify encoding for building on Windows
-        with open(doc_path, "w+", encoding="utf-8") as f:
-            f.write(example.generate())
-        logger.debug("Example generated: %s", doc_path.relative_to(ROOT_DIR))
-
-    # Update the navigation file
-    update_nav_file(examples)

docs/mkdocs/stylesheets/extra.css

@@ -25,11 +25,8 @@ a:not(:has(svg)):not(.md-icon):not(.autorefs-external) {
 
 a[href*="localhost"]::after,
 a[href*="127.0.0.1"]::after,
-
-/* Hide external link icons for all links */
-a[href^="//"]::after,
-a[href^="http://"]::after,
-a[href^="https://"]::after {
+a[href*="org.readthedocs.build"]::after,
+a[href*="docs.vllm.ai"]::after {
     display: none !important;
 }
 

docs/models/supported_models.md

@@ -4,7 +4,7 @@ vLLM-Omni supports unified multimodal comprehension and generation models across
 
 ## Model Implementation
 
-If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
+If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm-omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
 
 ## List of Supported Models for Nvidia GPU
 
@@ -17,14 +17,11 @@ th {
 
 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
-| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
-| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
-|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
 
 
-## List of Supported Models for NPU
+## List of Supported Models for Ascend NPU
 
 <style>
 th {
@@ -36,4 +33,3 @@ th {
 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
-| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |

docs/usage/faq.md

@@ -2,7 +2,7 @@
 
 > 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 have one AutoRegressive model stage. This is because the unified KV cache management of vLLM. Stages of other types can coexist within a chip. The restriction will be resolved in later version.
+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?
 
@@ -11,15 +11,3 @@ A: If you encounter error about backend of librosa, try to install ffmpeg with c
 sudo apt update
 sudo apt install ffmpeg
 ```
-
-> Q: I encounter some bugs or CI problems, which is urgent. How can I solve it?
-
-A: At first, you can check current [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions. If none of these satisfy your demand and it is urgent, please find these [volunteers](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/volunteers/) for help.
-
-> Q: Does vLLM-Omni support AWQ or any other quantization?
-
-A: vLLM-Omni partitions model into several stages. For AR stages, it will reuse main logic of LLMEngine in vLLM. So current quantization supported in vLLM should be also supported in vLLM-Omni for them. But systematic verification is ongoing. For quantization for DiffusionEngine, we are working on it. Please stay tuned and welcome contribution!
-
-> Q: Does vLLM-Omni support multimodal streaming input and output?
-
-A: Not yet. We already put it on the [Roadmap](https://github.com/vllm-project/vllm-omni/issues/165). Please stay tuned!

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -1,70 +1,55 @@
-# Offline Example of vLLM-Omni for Qwen2.5-Omni
+# Offline Inference Example of vLLM-Omni for Qwen2.5-Omni
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni>.
 
 
 ## 🛠️ Installation
 
-Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
+Please refer to [installation](../../../getting_started/installation/README.md).
 
-## Run examples (Qwen2.5-Omni)
+## Run Offline inference with Qwen2.5-Omni.
+First, navigate to the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Inside the directory, `end2end.py` is a comprehensive demo suite for initializing a model instance of `Qwen/Qwen2.5-Omni-7B` and use it for various offline inference tasks.
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen2_5_omni/end2end.py"
+    ``````
+
+Below we also provide simple bash scripts to execute this file.
+### Single Prompt
+```bash
+bash run_single_prompt.sh
+```
+??? abstract "run_single_prompt.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
+    ``````
 
 ### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit) and extract the prompts with `extract_prompts.py`
+??? abstract "extract_prompts.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
+    ``````
 ```bash
 tar -xf <Your Download Path>/seedtts_testset.tar
 cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
 python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
   --input examples/offline_inference/qwen2_5_omni/meta.lst \
   --output examples/offline_inference/qwen2_5_omni/top100.txt \
-  --topk 100
-```
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
+  --topk 10
 ```
 Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
-
-### Single Prompt
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Then run the command below.
-```bash
-bash run_single_prompt.sh
-```
-
-#### Using Local Media Files
-The `end2end.py` script supports local media files (audio, video, image) via CLI arguments:
-
-```bash
-# Use single local media files
-python end2end.py --query-type use_image --image-path /path/to/image.jpg
-python end2end.py --query-type use_video --video-path /path/to/video.mp4
-python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
-
-# Combine multiple local media files
-python end2end.py --query-type mixed_modalities \
-    --video-path /path/to/video.mp4 \
-    --image-path /path/to/image.jpg \
-    --audio-path /path/to/audio.wav
-
-# Use audio from video file
-python end2end.py --query-type use_audio_in_video --video-path /path/to/video.mp4
-
-```
-
-If media file paths are not provided, the script will use default assets. Supported query types:
-- `use_image`: Image input only
-- `use_video`: Video input only
-- `use_audio`: Audio input only
-- `mixed_modalities`: Audio + image + video
-- `use_audio_in_video`: Extract audio from video
-- `text`: Text-only query
+??? abstract "run_multiple_prompts.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
+    ``````
 
 ### FAQ
 
@@ -73,22 +58,3 @@ If you encounter error about backend of librosa, try to install ffmpeg with comm
 sudo apt update
 sudo apt install ffmpeg
 ```
-
-## Example materials
-
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen2_5_omni/end2end.py"
-    ``````
-??? abstract "extract_prompts.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
-    ``````
-??? abstract "run_multiple_prompts.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
-    ``````
-??? abstract "run_single_prompt.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
-    ``````

docs/user_guide/examples/offline_inference/qwen3_omni.md

@@ -1,81 +0,0 @@
-# Offline Example of vLLM-Omni for Qwen3-Omni
-
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_omni>.
-
-
-## 🛠️ Installation
-
-Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
-
-## Run examples (Qwen3-Omni)
-
-### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-Omni README.md](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni/README.md)
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen3_omni
-```
-Then run the command below.
-```bash
-bash run_multiple_prompts.sh
-```
-### Single Prompt
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen3_omni
-```
-Then run the command below.
-```bash
-bash run_single_prompt.sh
-```
-If you have not enough memory, you can set thinker with tensor parallel. Just run the command below.
-```bash
-bash run_single_prompt_tp.sh
-```
-
-#### Using Local Media Files
-The `end2end.py` script supports local media files (audio, video, image) via command-line arguments:
-
-```bash
-# Use local video file
-python end2end.py --query-type use_video --video-path /path/to/video.mp4
-
-# Use local image file
-python end2end.py --query-type use_image --image-path /path/to/image.jpg
-
-# Use local audio file
-python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
-```
-
-If media file paths are not provided, the script will use default assets. Supported query types:
-- `use_video`: Video input
-- `use_image`: Image input
-- `use_audio`: Audio input
-- `text`: Text-only query
-
-### 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/qwen3_omni/end2end.py"
-    ``````
-??? abstract "run_multiple_prompts.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
-    ``````
-??? abstract "run_single_prompt.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt.sh"
-    ``````
-??? abstract "run_single_prompt_tp.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
-    ``````

docs/user_guide/examples/offline_inference/qwen_image.md

@@ -1,57 +0,0 @@
-# Qwen-Image Offline Inference
-
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen_image>.
-
-
-This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Image` using vLLM-Omni:
-
-- `text_to_image.py`: command-line script for single image generation.
-- `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
-
-
-## Local CLI Usage
-
-```bash
-python text_to_image.py \
-  --prompt "a cup of coffee on the table" \
-  --seed 42 \
-  --cfg_scale 4.0 \
-  --num_images_per_prompt 1 \
-  --num_inference_steps 50 \
-  --height 1024 \
-  --width 1024 \
-  --output outputs/coffee.png
-```
-
-Key arguments:
-
-- `--prompt`: text description (string).
-- `--seed`: integer seed for deterministic sampling.
-- `--cfg_scale`: true CFG scale (model-specific guidance strength).
-- `--num_images_per_prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
-- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
-- `--height/--width`: output resolution (defaults 1024x1024).
-- `--output`: path to save the generated PNG.
-
-> ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
-
-## Web UI Demo
-
-Launch the gradio demo:
-
-```bash
-python gradio_demo.py --port 7862
-```
-
-Then open `http://localhost:7862/` on your local browser to interact with the web UI.
-
-## Example materials
-
-??? abstract "gradio_demo.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen_image/gradio_demo.py"
-    ``````
-??? abstract "text_to_image.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen_image/text_to_image.py"
-    ``````

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -1,99 +1,43 @@
-# Online serving Example of vLLM-Omni for Qwen2.5-Omni
+# Online Serving Example of vLLM-Omni for Qwen2.5-Omni
 
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen2_5_omni>.
+Source <https://github.com/vllm-project/vllm-omni/blob/main/examples/online_serving/README.md>.
 
 
 ## 🛠️ Installation
 
-Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
+Please refer to [installation](../../../getting_started/installation/README.md).
 
-## Run examples (Qwen2.5-Omni)
-
-### Launch the Server
+## Deploy Qwen/Qwen2.5-Omni-7B
 
+First launch the OpenAI-compatible inference server
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
-
 If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
-
-### Send Multi-modal Request
-
-Get into the example folder
+## Query the model
+Navigate to the example folder
 ```bash
-cd examples/online_serving/qwen2_5_omni
+cd examples/online_serving
 ```
-
-####  Send request via python
-
+Query the model server with OpenAI Python API client:
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```
-
-The Python client supports the following command-line arguments:
-
-- `--query-type` (or `-q`): Query type (default: `mixed_modalities`)
-  - Options: `mixed_modalities`, `use_audio_in_video`, `multi_audios`, `text`
-- `--video-path` (or `-v`): Path to local video file or URL
-  - If not provided and query-type uses video, uses default video URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
-- `--image-path` (or `-i`): Path to local image file or URL
-  - If not provided and query-type uses image, uses default image URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common image formats: JPEG, PNG, GIF, WebP
-  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
-- `--audio-path` (or `-a`): Path to local audio file or URL
-  - If not provided and query-type uses audio, uses default audio URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
-  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
-- `--prompt` (or `-p`): Custom text prompt/question
-  - If not provided, uses default prompt for the selected query type
-  - Example: `--prompt "What are the main activities shown in this video?"`
-
-
-For example, to use mixed modalities with all local files:
-
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py \
-    --query-type mixed_modalities \
-    --video-path /path/to/your/video.mp4 \
-    --image-path /path/to/your/image.jpg \
-    --audio-path /path/to/your/audio.wav \
-    --prompt "Analyze all the media content and provide a comprehensive summary."
-```
-
-####  Send request via curl
-
+??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
+    ``````py
+    --8<-- "examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py"
+    ``````
+You can also query the model with `curl` command:
 ```bash
 bash run_curl_multimodal_generation.sh mixed_modalities
 ```
-
-## 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 AsyncOmni 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.
-
-The gradio script supports the following arguments:
-
-- `--model`: Model name
-- `--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)
+??? abstract "run_curl_multimodal_generation.sh"
+    ``````py
+    --8<-- "examples/online_serving/run_curl_multimodal_generation.sh"
+    ``````
 
 ### FAQ
 
@@ -103,17 +47,34 @@ sudo apt update
 sudo apt install ffmpeg
 ```
 
-## Example materials
+## Run Local Web UI Demo
 
-??? abstract "gradio_demo.py"
-    ``````py
-    --8<-- "examples/online_serving/qwen2_5_omni/gradio_demo.py"
-    ``````
-??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
-    ``````py
-    --8<-- "examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py"
-    ``````
-??? abstract "run_curl_multimodal_generation.sh"
-    ``````sh
-    --8<-- "examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh"
-    ``````
+You can also deploy a Gradio Web UI that allows users to interact with the model through a web browser. Below is an example on how to do so with `Qwen/Qwen2.5-Omni-7B`.
+
+### Running Gradio Demo
+
+Install gradio with `uv pip install "gradio>=5.49.1,<6.0.0"`, then you can launch the web service built on AsyncOmni by
+
+```bash
+python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861
+```
+Now you can interact with model via the web UI at `http://localhost:7861/` on your local browser.
+
+
+### Options
+
+You can customize its basic launch parameters:
+
+```bash
+python gradio_demo.py \
+    --model Qwen/Qwen2.5-Omni-7B \
+    --ip 127.0.0.1 \
+    --port 7861 \
+    --stage-configs-path /path/to/stage_configs.yaml
+```
+
+- `--model`: Local model checkpoint to load (default `Qwen/Qwen2.5-Omni-7B`).
+- `--ip`: Host/IP for the Gradio server (default `127.0.0.1`).
+- `--port`: Port for the Gradio server (default `7861`).
+- `--stage-configs-path`: Optional path to custom stage configs YAML.
+- `--share`: Set to expose a temporary public link via Gradio.

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -1,124 +0,0 @@
-# Online serving Example of vLLM-Omni for Qwen3-omni
-
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni>.
-
-
-## 🛠️ Installation
-
-Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
-
-## Run examples (Qwen3-Omni)
-
-### Launch the Server
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
-```
-
-If you have custom stage configs file, launch the server with command below
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
-```
-
-### Send Multi-modal Request
-
-Get into the example folder
-```bash
-cd examples/online_serving/qwen3_omni
-```
-
-####  Send request via python
-
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
-```
-
-The Python client supports the following command-line arguments:
-
-- `--query-type` (or `-q`): Query type (default: `use_video`)
-  - Options: `text`, `use_audio`, `use_image`, `use_video`
-- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-Omni-30B-A3B-Instruct`)
-- `--video-path` (or `-v`): Path to local video file or URL
-  - If not provided and query-type is `use_video`, uses default video URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
-- `--image-path` (or `-i`): Path to local image file or URL
-  - If not provided and query-type is `use_image`, uses default image URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common image formats: JPEG, PNG, GIF, WebP
-  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
-- `--audio-path` (or `-a`): Path to local audio file or URL
-  - If not provided and query-type is `use_audio`, uses default audio URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
-  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
-- `--prompt` (or `-p`): Custom text prompt/question
-  - If not provided, uses default prompt for the selected query type
-  - Example: `--prompt "What are the main activities shown in this video?"`
-
-
-For example, to use a local video file with custom prompt:
-
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py \
-    --query-type use_video \
-    --video-path /path/to/your/video.mp4 \
-    --prompt "What are the main activities shown in this video?"
-```
-
-####  Send request via curl
-
-```bash
-bash run_curl_multimodal_generation.sh use_image
-```
-
-
-### 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 AsyncOmni by
-
-```bash
-python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
-```
-
-Then open `http://localhost:7861/` on your local browser to interact with the web UI.
-
-
-The gradio script supports the following arguments:
-
-- `--model`: Model name
-- `--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)
-
-## Example materials
-
-??? abstract "gradio_demo.py"
-    ``````py
-    --8<-- "examples/online_serving/qwen3_omni/gradio_demo.py"
-    ``````
-??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
-    ``````py
-    --8<-- "examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py"
-    ``````
-??? abstract "qwen3_omni_moe_thinking.yaml"
-    ``````yaml
-    --8<-- "examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml"
-    ``````
-??? abstract "run_curl_multimodal_generation.sh"
-    ``````sh
-    --8<-- "examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh"
-    ``````

examples/offline_inference/qwen2_5_omni/README.md

@@ -1,11 +1,10 @@
-# Offline Example of vLLM-Omni for Qwen2.5-Omni
+# Offline Example of vLLM-Omni for Qwen2.5-omni
 
 ## 🛠️ Installation
 
 Please refer to [README.md](../../../README.md)
 
-## Run examples (Qwen2.5-Omni)
-
+## Run examples (Qwen2.5-omni)
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
 ```bash
@@ -24,8 +23,7 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
-
-### Single Prompt
+### Single Prompts
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen2_5_omni
@@ -35,34 +33,6 @@ Then run the command below.
 bash run_single_prompt.sh
 ```
 
-#### Using Local Media Files
-The `end2end.py` script supports local media files (audio, video, image) via CLI arguments:
-
-```bash
-# Use single local media files
-python end2end.py --query-type use_image --image-path /path/to/image.jpg
-python end2end.py --query-type use_video --video-path /path/to/video.mp4
-python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
-
-# Combine multiple local media files
-python end2end.py --query-type mixed_modalities \
-    --video-path /path/to/video.mp4 \
-    --image-path /path/to/image.jpg \
-    --audio-path /path/to/audio.wav
-
-# Use audio from video file
-python end2end.py --query-type use_audio_in_video --video-path /path/to/video.mp4
-
-```
-
-If media file paths are not provided, the script will use default assets. Supported query types:
-- `use_image`: Image input only
-- `use_video`: Video input only
-- `use_audio`: Audio input only
-- `mixed_modalities`: Audio + image + video
-- `use_audio_in_video`: Extract audio from video
-- `text`: Text-only query
-
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -6,20 +6,17 @@ with the correct prompt format on Qwen2.5-Omni
 """
 
 import os
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 
-import librosa
-import numpy as np
 import soundfile as sf
-from PIL import Image
 from vllm.assets.audio import AudioAsset
 from vllm.assets.image import ImageAsset
-from vllm.assets.video import VideoAsset, video_to_ndarrays
+from vllm.assets.video import VideoAsset
 from vllm.multimodal.image import convert_image_mode
 from vllm.sampling_params import SamplingParams
 from vllm.utils import FlexibleArgumentParser
 
-from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.entrypoints.omni_llm import OmniLLM
 
 SEED = 42
 
@@ -57,13 +54,7 @@ def get_text_query(question: str = None) -> QueryResult:
     )
 
 
-def get_mixed_modalities_query(
-    video_path: Optional[str] = None,
-    image_path: Optional[str] = None,
-    audio_path: Optional[str] = None,
-    num_frames: int = 16,
-    sampling_rate: int = 16000,
-) -> QueryResult:
+def get_mixed_modalities_query() -> QueryResult:
     question = "What is recited in the audio? What is the content of this image? Why is this video funny?"
     prompt = (
         f"<|im_start|>system\n{default_system}<|im_end|>\n"
@@ -73,49 +64,20 @@ def get_mixed_modalities_query(
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
-
-    # Load video
-    if video_path:
-        if not os.path.exists(video_path):
-            raise FileNotFoundError(f"Video file not found: {video_path}")
-        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
-    else:
-        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
-
-    # Load image
-    if image_path:
-        if not os.path.exists(image_path):
-            raise FileNotFoundError(f"Image file not found: {image_path}")
-        pil_image = Image.open(image_path)
-        image_data = convert_image_mode(pil_image, "RGB")
-    else:
-        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
-
-    # Load audio
-    if audio_path:
-        if not os.path.exists(audio_path):
-            raise FileNotFoundError(f"Audio file not found: {audio_path}")
-        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
-        audio_data = (audio_signal.astype(np.float32), sr)
-    else:
-        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
-
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "audio": audio_data,
-                "image": image_data,
-                "video": video_frames,
+                "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
+                "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
+                "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
             },
         },
         limit_mm_per_prompt={"audio": 1, "image": 1, "video": 1},
     )
 
 
-def get_use_audio_in_video_query(
-    video_path: Optional[str] = None, num_frames: int = 16, sampling_rate: int = 16000
-) -> QueryResult:
+def get_use_audio_in_video_query() -> QueryResult:
     question = "Describe the content of the video, then convert what the baby say into text."
     prompt = (
         f"<|im_start|>system\n{default_system}<|im_end|>\n"
@@ -123,24 +85,14 @@ def get_use_audio_in_video_query(
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
-
-    if video_path:
-        if not os.path.exists(video_path):
-            raise FileNotFoundError(f"Video file not found: {video_path}")
-        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
-        # Extract audio from video file
-        audio_signal, sr = librosa.load(video_path, sr=sampling_rate)
-        audio = (audio_signal.astype(np.float32), sr)
-    else:
-        asset = VideoAsset(name="baby_reading", num_frames=num_frames)
-        video_frames = asset.np_ndarrays
-        audio = asset.get_audio(sampling_rate=sampling_rate)
+    asset = VideoAsset(name="baby_reading", num_frames=16)
+    audio = asset.get_audio(sampling_rate=16000)
 
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "video": video_frames,
+                "video": asset.np_ndarrays,
                 "audio": audio,
             },
             "mm_processor_kwargs": {
@@ -151,7 +103,7 @@ def get_use_audio_in_video_query(
     )
 
 
-def get_multi_audios_query(audio_path: Optional[str] = None, sampling_rate: int = 16000) -> QueryResult:
+def get_multi_audios_query() -> QueryResult:
     question = "Are these two audio clips the same?"
     prompt = (
         f"<|im_start|>system\n{default_system}<|im_end|>\n"
@@ -160,28 +112,14 @@ def get_multi_audios_query(audio_path: Optional[str] = None, sampling_rate: int
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
-
-    if audio_path:
-        if not os.path.exists(audio_path):
-            raise FileNotFoundError(f"Audio file not found: {audio_path}")
-        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
-        audio_data = (audio_signal.astype(np.float32), sr)
-        # Use the provided audio as the first audio, default as second
-        audio_list = [
-            audio_data,
-            AudioAsset("mary_had_lamb").audio_and_sample_rate,
-        ]
-    else:
-        audio_list = [
-            AudioAsset("winning_call").audio_and_sample_rate,
-            AudioAsset("mary_had_lamb").audio_and_sample_rate,
-        ]
-
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "audio": audio_list,
+                "audio": [
+                    AudioAsset("winning_call").audio_and_sample_rate,
+                    AudioAsset("mary_had_lamb").audio_and_sample_rate,
+                ],
             },
         },
         limit_mm_per_prompt={
@@ -190,137 +128,19 @@ def get_multi_audios_query(audio_path: Optional[str] = None, sampling_rate: int
     )
 
 
-def get_image_query(question: str = None, image_path: Optional[str] = None) -> QueryResult:
-    if question is None:
-        question = "What is the content of this image?"
-    prompt = (
-        f"<|im_start|>system\n{default_system}<|im_end|>\n"
-        "<|im_start|>user\n<|vision_bos|><|IMAGE|><|vision_eos|>"
-        f"{question}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
-    )
-
-    if image_path:
-        if not os.path.exists(image_path):
-            raise FileNotFoundError(f"Image file not found: {image_path}")
-        pil_image = Image.open(image_path)
-        image_data = convert_image_mode(pil_image, "RGB")
-    else:
-        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
-
-    return QueryResult(
-        inputs={
-            "prompt": prompt,
-            "multi_modal_data": {
-                "image": image_data,
-            },
-        },
-        limit_mm_per_prompt={"image": 1},
-    )
-
-
-def get_video_query(question: str = None, video_path: Optional[str] = None, num_frames: int = 16) -> QueryResult:
-    if question is None:
-        question = "Why is this video funny?"
-    prompt = (
-        f"<|im_start|>system\n{default_system}<|im_end|>\n"
-        "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|>"
-        f"{question}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
-    )
-
-    if video_path:
-        if not os.path.exists(video_path):
-            raise FileNotFoundError(f"Video file not found: {video_path}")
-        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
-    else:
-        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
-
-    return QueryResult(
-        inputs={
-            "prompt": prompt,
-            "multi_modal_data": {
-                "video": video_frames,
-            },
-        },
-        limit_mm_per_prompt={"video": 1},
-    )
-
-
-def get_audio_query(question: str = None, audio_path: Optional[str] = None, sampling_rate: int = 16000) -> QueryResult:
-    if question is None:
-        question = "What is the content of this audio?"
-    prompt = (
-        f"<|im_start|>system\n{default_system}<|im_end|>\n"
-        "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
-        f"{question}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
-    )
-
-    if audio_path:
-        if not os.path.exists(audio_path):
-            raise FileNotFoundError(f"Audio file not found: {audio_path}")
-        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
-        audio_data = (audio_signal.astype(np.float32), sr)
-    else:
-        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
-
-    return QueryResult(
-        inputs={
-            "prompt": prompt,
-            "multi_modal_data": {
-                "audio": audio_data,
-            },
-        },
-        limit_mm_per_prompt={"audio": 1},
-    )
-
-
 query_map = {
     "mixed_modalities": get_mixed_modalities_query,
     "use_audio_in_video": get_use_audio_in_video_query,
     "multi_audios": get_multi_audios_query,
-    "use_image": get_image_query,
-    "use_video": get_video_query,
-    "use_audio": get_audio_query,
     "text": get_text_query,
 }
 
 
 def main(args):
     model_name = "Qwen/Qwen2.5-Omni-7B"
+    query_result = query_map[args.query_type]()
 
-    # Get paths from args
-    video_path = getattr(args, "video_path", None)
-    image_path = getattr(args, "image_path", None)
-    audio_path = getattr(args, "audio_path", None)
-    num_frames = getattr(args, "num_frames", 16)
-    sampling_rate = getattr(args, "sampling_rate", 16000)
-
-    # Get the query function and call it with appropriate parameters
-    query_func = query_map[args.query_type]
-    if args.query_type == "mixed_modalities":
-        query_result = query_func(
-            video_path=video_path,
-            image_path=image_path,
-            audio_path=audio_path,
-            num_frames=num_frames,
-            sampling_rate=sampling_rate,
-        )
-    elif args.query_type == "use_audio_in_video":
-        query_result = query_func(video_path=video_path, num_frames=num_frames, sampling_rate=sampling_rate)
-    elif args.query_type == "multi_audios":
-        query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
-    elif args.query_type == "use_image":
-        query_result = query_func(image_path=image_path)
-    elif args.query_type == "use_video":
-        query_result = query_func(video_path=video_path, num_frames=num_frames)
-    elif args.query_type == "use_audio":
-        query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
-    else:
-        query_result = query_func()
-
-    omni_llm = Omni(
+    omni_llm = OmniLLM(
         model=model_name,
         log_stats=args.enable_stats,
         log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
@@ -463,39 +283,6 @@ def parse_args():
         default=None,
         help="Path to a .txt file with one prompt per line (preferred).",
     )
-    parser.add_argument(
-        "--video-path",
-        "-v",
-        type=str,
-        default=None,
-        help="Path to local video file. If not provided, uses default video asset.",
-    )
-    parser.add_argument(
-        "--image-path",
-        "-i",
-        type=str,
-        default=None,
-        help="Path to local image file. If not provided, uses default image asset.",
-    )
-    parser.add_argument(
-        "--audio-path",
-        "-a",
-        type=str,
-        default=None,
-        help="Path to local audio file. If not provided, uses default audio asset.",
-    )
-    parser.add_argument(
-        "--num-frames",
-        type=int,
-        default=16,
-        help="Number of frames to extract from video (default: 16).",
-    )
-    parser.add_argument(
-        "--sampling-rate",
-        type=int,
-        default=16000,
-        help="Sampling rate for audio loading (default: 16000).",
-    )
 
     return parser.parse_args()
 

examples/offline_inference/qwen2_5_omni/extract_prompts.py

@@ -4,7 +4,7 @@ from typing import Optional
 
 
 def extract_prompt(line: str) -> Optional[str]:
-    # Extract the content between the first '|' and the second '|'
+    # 提取第一个 '|' 与第二个 '|' 之间的内容
     i = line.find("|")
     if i == -1:
         return None

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

@@ -1,2 +1,2 @@
 python end2end.py --output-wav output_audio \
-                  --query-type mixed_modalities
+                  --query-type use_audio_in_video

examples/offline_inference/qwen3_omni/README.md

@@ -1,13 +1,12 @@
-# Offline Example of vLLM-Omni for Qwen3-Omni
+# Offline Example of vLLM-Omni for Qwen2.5-omni
 
 ## 🛠️ Installation
 
 Please refer to [README.md](../../../README.md)
 
-## Run examples (Qwen3-Omni)
-
+## Run examples (Qwen3-omni)
 ### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-Omni README.md](../qwen2_5_omni/README.md)
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-omni README.md](../qwen2_5_omni/README.md)
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -16,7 +15,7 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
-### Single Prompt
+### Single Prompts
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -25,30 +24,6 @@ Then run the command below.
 ```bash
 bash run_single_prompt.sh
 ```
-If you have not enough memory, you can set thinker with tensor parallel. Just run the command below.
-```bash
-bash run_single_prompt_tp.sh
-```
-
-#### Using Local Media Files
-The `end2end.py` script supports local media files (audio, video, image) via command-line arguments:
-
-```bash
-# Use local video file
-python end2end.py --query-type use_video --video-path /path/to/video.mp4
-
-# Use local image file
-python end2end.py --query-type use_image --image-path /path/to/image.jpg
-
-# Use local audio file
-python end2end.py --query-type use_audio --audio-path /path/to/audio.wav
-```
-
-If media file paths are not provided, the script will use default assets. Supported query types:
-- `use_video`: Video input
-- `use_image`: Image input
-- `use_audio`: Audio input
-- `text`: Text-only query
 
 ### FAQ
 

examples/offline_inference/qwen3_omni/end2end.py

@@ -6,20 +6,17 @@ with the correct prompt format on Qwen3-Omni (thinker only).
 """
 
 import os
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 
-import librosa
-import numpy as np
 import soundfile as sf
-from PIL import Image
 from vllm import SamplingParams
 from vllm.assets.audio import AudioAsset
 from vllm.assets.image import ImageAsset
-from vllm.assets.video import VideoAsset, video_to_ndarrays
+from vllm.assets.video import VideoAsset
 from vllm.multimodal.image import convert_image_mode
 from vllm.utils import FlexibleArgumentParser
 
-from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.entrypoints.omni_llm import OmniLLM
 
 SEED = 42
 
@@ -57,7 +54,7 @@ def get_text_query(question: str = None) -> QueryResult:
     )
 
 
-def get_video_query(question: str = None, video_path: Optional[str] = None, num_frames: int = 16) -> QueryResult:
+def get_video_query(question: str = None) -> QueryResult:
     if question is None:
         question = "Why is this video funny?"
     prompt = (
@@ -67,25 +64,18 @@ def get_video_query(question: str = None, video_path: Optional[str] = None, num_
         f"<|im_start|>assistant\n"
     )
 
-    if video_path:
-        if not os.path.exists(video_path):
-            raise FileNotFoundError(f"Video file not found: {video_path}")
-        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
-    else:
-        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
-
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "video": video_frames,
+                "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
             },
         },
         limit_mm_per_prompt={"video": 1},
     )
 
 
-def get_image_query(question: str = None, image_path: Optional[str] = None) -> QueryResult:
+def get_image_query(question: str = None) -> QueryResult:
     if question is None:
         question = "What is the content of this image?"
     prompt = (
@@ -94,27 +84,18 @@ def get_image_query(question: str = None, image_path: Optional[str] = None) -> Q
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
-
-    if image_path:
-        if not os.path.exists(image_path):
-            raise FileNotFoundError(f"Image file not found: {image_path}")
-        pil_image = Image.open(image_path)
-        image_data = convert_image_mode(pil_image, "RGB")
-    else:
-        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
-
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "image": image_data,
+                "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
             },
         },
         limit_mm_per_prompt={"image": 1},
     )
 
 
-def get_audio_query(question: str = None, audio_path: Optional[str] = None, sampling_rate: int = 16000) -> QueryResult:
+def get_audio_query(question: str = None) -> QueryResult:
     if question is None:
         question = "What is the content of this audio?"
     prompt = (
@@ -123,20 +104,11 @@ def get_audio_query(question: str = None, audio_path: Optional[str] = None, samp
         f"{question}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
-
-    if audio_path:
-        if not os.path.exists(audio_path):
-            raise FileNotFoundError(f"Audio file not found: {audio_path}")
-        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
-        audio_data = (audio_signal.astype(np.float32), sr)
-    else:
-        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
-
     return QueryResult(
         inputs={
             "prompt": prompt,
             "multi_modal_data": {
-                "audio": audio_data,
+                "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
             },
         },
         limit_mm_per_prompt={"audio": 1},
@@ -153,26 +125,10 @@ query_map = {
 
 def main(args):
     model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+    query_result = query_map[args.query_type]()
 
-    # Get paths from args
-    video_path = getattr(args, "video_path", None)
-    image_path = getattr(args, "image_path", None)
-    audio_path = getattr(args, "audio_path", None)
-
-    # Get the query function and call it with appropriate parameters
-    query_func = query_map[args.query_type]
-    if args.query_type == "use_video":
-        query_result = query_func(video_path=video_path, num_frames=getattr(args, "num_frames", 16))
-    elif args.query_type == "use_image":
-        query_result = query_func(image_path=image_path)
-    elif args.query_type == "use_audio":
-        query_result = query_func(audio_path=audio_path, sampling_rate=getattr(args, "sampling_rate", 16000))
-    else:
-        query_result = query_func()
-
-    omni_llm = Omni(
+    omni_llm = OmniLLM(
         model=model_name,
-        stage_configs_path=args.stage_configs_path,
     )
 
     thinker_sampling_params = SamplingParams(
@@ -320,45 +276,6 @@ def parse_args():
         default=None,
         help="Path to a .txt file with one prompt per line (preferred).",
     )
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=None,
-        help="Path to a stage configs file.",
-    )
-    parser.add_argument(
-        "--video-path",
-        "-v",
-        type=str,
-        default=None,
-        help="Path to local video file. If not provided, uses default video asset.",
-    )
-    parser.add_argument(
-        "--image-path",
-        "-i",
-        type=str,
-        default=None,
-        help="Path to local image file. If not provided, uses default image asset.",
-    )
-    parser.add_argument(
-        "--audio-path",
-        "-a",
-        type=str,
-        default=None,
-        help="Path to local audio file. If not provided, uses default audio asset.",
-    )
-    parser.add_argument(
-        "--num-frames",
-        type=int,
-        default=16,
-        help="Number of frames to extract from video (default: 16).",
-    )
-    parser.add_argument(
-        "--sampling-rate",
-        type=int,
-        default=16000,
-        help="Sampling rate for audio loading (default: 16000).",
-    )
 
     return parser.parse_args()
 
@@ -366,3 +283,7 @@ def parse_args():
 if __name__ == "__main__":
     args = parse_args()
     main(args)
+
+    # use examples:
+    # python end2end.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type text
+    # python end2end.py --model /custom_path/Qwen3-Omni-30B-A3B-Instruct --query-type use_video_only

examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh

@@ -1,5 +0,0 @@
-python end2end.py --output-wav output_audio \
-                  --query-type use_audio \
-                  --init-sleep-seconds 90
-
-# init-sleep-seconds works to avoid two vLLM stages initialized at the same time within a card.

examples/offline_inference/qwen_image/README.md

@@ -1,21 +1,10 @@
 # Qwen-Image Offline Inference
 
-This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` using vLLM-Omni:
+This folder provides two simple entrypoints for experimenting with `Qwen/Qwen-Image` using vLLM-Omni:
 
-- `text_to_image.py`: command-line script for single image generation with advanced options.
+- `text_to_image.py`: command-line script for single image generation.
 - `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
 
-## Basic Usage
-
-```python
-from vllm_omni.entrypoints.omni import Omni
-
-if __name__ == "__main__":
-    omni = Omni(model="Qwen/Qwen-Image")
-    prompt = "a cup of coffee on the table"
-    images = omni.generate(prompt)
-    images[0].save("coffee.png")
-```
 
 ## Local CLI Usage
 

examples/offline_inference/qwen_image/gradio_demo.py

@@ -5,7 +5,6 @@ import gradio as gr
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
     "1:1": (1328, 1328),
@@ -59,18 +58,11 @@ def parse_args() -> argparse.Namespace:
 
 @lru_cache(maxsize=1)
 def get_omni(model_name: str) -> Omni:
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
-    return Omni(
-        model=model_name,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
-    )
+    return Omni(model=model_name)
 
 
 def build_demo(args: argparse.Namespace) -> gr.Blocks:
-    device = detect_device_type()
+    device = "cuda"
     omni = get_omni(args.model)
 
     def run_inference(
@@ -105,6 +97,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
             generator=generator,
             true_cfg_scale=float(cfg_scale_value),
             num_inference_steps=num_steps,
+            num_images_per_prompt=num_images,
             num_outputs_per_prompt=num_images,
         )
         return [img for img in images[:num_images]]

examples/offline_inference/qwen_image/image_edit.py

@@ -1,143 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Example script for image editing with Qwen-Image-Edit.
-
-Usage:
-    python image_edit.py \
-        --image input.png \
-        --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
-        --output output_image_edit.png \
-        --num_inference_steps 50 \
-        --cfg_scale 4.0
-
-For more options, run:
-    python image_edit.py --help
-"""
-
-import argparse
-import os
-from pathlib import Path
-
-import torch
-from PIL import Image
-
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.")
-    parser.add_argument(
-        "--model",
-        default="Qwen/Qwen-Image-Edit",
-        help="Diffusion model name or local path.",
-    )
-    parser.add_argument(
-        "--image",
-        type=str,
-        required=True,
-        help="Path to input image file (PNG, JPG, etc.).",
-    )
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        required=True,
-        help="Text prompt describing the edit to make to the image.",
-    )
-    parser.add_argument(
-        "--negative_prompt",
-        type=str,
-        default=" ",
-        required=False,
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=0,
-        help="Random seed for deterministic results.",
-    )
-    parser.add_argument(
-        "--cfg_scale",
-        type=float,
-        default=4.0,
-        help="True classifier-free guidance scale specific to Qwen-Image-Edit.",
-    )
-    parser.add_argument(
-        "--output",
-        type=str,
-        default="output_image_edit.png",
-        help="Path to save the edited image (PNG).",
-    )
-    parser.add_argument(
-        "--num_outputs_per_prompt",
-        type=int,
-        default=1,
-        help="Number of images to generate for the given prompt.",
-    )
-    parser.add_argument(
-        "--num_inference_steps",
-        type=int,
-        default=50,
-        help="Number of denoising steps for the diffusion sampler.",
-    )
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-
-    # Validate input image exists
-    if not os.path.exists(args.image):
-        raise FileNotFoundError(f"Input image not found: {args.image}")
-
-    # Load input image
-    input_image = Image.open(args.image).convert("RGB")
-    print(f"Loaded input image from {args.image} (size: {input_image.size})")
-
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
-
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
-
-    # Initialize Omni with QwenImageEditPipeline
-    omni = Omni(
-        model=args.model,
-        model_class_name="QwenImageEditPipeline",
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
-    )
-    print("Pipeline loaded")
-
-    # Generate edited image
-    images = omni.generate(
-        prompt=args.prompt,
-        pil_image=input_image,
-        negative_prompt=args.negative_prompt,
-        generator=generator,
-        true_cfg_scale=args.cfg_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_outputs_per_prompt=args.num_outputs_per_prompt,
-    )
-
-    # Save output image(s)
-    output_path = Path(args.output)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    suffix = output_path.suffix or ".png"
-    stem = output_path.stem or "output_image_edit"
-
-    if args.num_outputs_per_prompt <= 1:
-        images[0].save(output_path)
-        print(f"Saved edited image to {os.path.abspath(output_path)}")
-    else:
-        for idx, img in enumerate(images):
-            save_path = output_path.parent / f"{stem}_{idx}{suffix}"
-            img.save(save_path)
-            print(f"Saved edited image to {os.path.abspath(save_path)}")
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/qwen_image/text_to_image.py

@@ -7,7 +7,6 @@ from pathlib import Path
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
 
 def parse_args() -> argparse.Namespace:
@@ -46,18 +45,10 @@ def parse_args() -> argparse.Namespace:
 
 def main():
     args = parse_args()
-    device = detect_device_type()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     generator = torch.Generator(device=device).manual_seed(args.seed)
 
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
-
-    omni = Omni(
-        model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
-    )
+    omni = Omni(model=args.model)
     images = omni.generate(
         args.prompt,
         height=args.height,
@@ -65,6 +56,7 @@ def main():
         generator=generator,
         true_cfg_scale=args.cfg_scale,
         num_inference_steps=args.num_inference_steps,
+        num_images_per_prompt=args.num_images_per_prompt,
         num_outputs_per_prompt=args.num_images_per_prompt,
     )
 

examples/online_serving/qwen2_5_omni/README.md

@@ -1,13 +1,12 @@
-# Online serving Example of vLLM-Omni for Qwen2.5-Omni
+# Online serving Example of vLLM-Omni for Qwen2.5-omni
 
 ## 🛠️ Installation
 
-Please refer to [README.md](../../../README.md)
+Please refer to [README.md](../../README.md)
 
-## Run examples (Qwen2.5-Omni)
-
-### Launch the Server
+## Run examples (Qwen2.5-omni)
 
+Launch the server
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
@@ -17,81 +16,21 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 
-### Send Multi-modal Request
-
 Get into the example folder
 ```bash
-cd examples/online_serving/qwen2_5_omni
+cd examples/online_serving
 ```
 
-####  Send request via python
-
+Send request via python
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```
 
-The Python client supports the following command-line arguments:
-
-- `--query-type` (or `-q`): Query type (default: `mixed_modalities`)
-  - Options: `mixed_modalities`, `use_audio_in_video`, `multi_audios`, `text`
-- `--video-path` (or `-v`): Path to local video file or URL
-  - If not provided and query-type uses video, uses default video URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
-- `--image-path` (or `-i`): Path to local image file or URL
-  - If not provided and query-type uses image, uses default image URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common image formats: JPEG, PNG, GIF, WebP
-  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
-- `--audio-path` (or `-a`): Path to local audio file or URL
-  - If not provided and query-type uses audio, uses default audio URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
-  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
-- `--prompt` (or `-p`): Custom text prompt/question
-  - If not provided, uses default prompt for the selected query type
-  - Example: `--prompt "What are the main activities shown in this video?"`
-
-
-For example, to use mixed modalities with all local files:
-
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py \
-    --query-type mixed_modalities \
-    --video-path /path/to/your/video.mp4 \
-    --image-path /path/to/your/image.jpg \
-    --audio-path /path/to/your/audio.wav \
-    --prompt "Analyze all the media content and provide a comprehensive summary."
-```
-
-####  Send request via curl
-
+Send request via curl
 ```bash
 bash run_curl_multimodal_generation.sh mixed_modalities
 ```
 
-## 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 AsyncOmni 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.
-
-The gradio script supports the following arguments:
-
-- `--model`: Model name
-- `--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)
-
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -99,3 +38,36 @@ If you encounter error about backend of librosa, try to install ffmpeg with comm
 sudo apt update
 sudo apt install ffmpeg
 ```
+
+## Run Local Web UI Demo
+
+This Web UI demo allows users to interact with the model through a web browser.
+
+### Running Gradio Demo
+
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+
+```bash
+python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861
+```
+
+Then open `http://localhost:7861/` on your local browser to interact with the web UI.
+
+
+### Options
+
+You can customize its basic launch parameters:
+
+```bash
+python gradio_demo.py \
+    --model Qwen/Qwen2.5-Omni-7B \
+    --ip 127.0.0.1 \
+    --port 7861 \
+    --stage-configs-path /path/to/stage_configs.yaml
+```
+
+- `--model`: Local model checkpoint to load (default `Qwen/Qwen2.5-Omni-7B`).
+- `--ip`: Host/IP for the Gradio server (default `127.0.0.1`).
+- `--port`: Port for the Gradio server (default `7861`).
+- `--stage-configs-path`: Optional path to custom stage configs YAML.
+- `--share`: Set to expose a temporary public link via Gradio.

examples/online_serving/qwen2_5_omni/gradio_demo.py

@@ -15,7 +15,7 @@ from PIL import Image
 from vllm.assets.video import video_get_metadata, video_to_ndarrays
 from vllm.sampling_params import SamplingParams
 
-from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
 
 # Import utils from offline inference example
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen2_5_omni"))
@@ -236,8 +236,8 @@ def process_video_file(
     return frames, metadata, audio_tuple
 
 
-async def run_inference_async_omni(
-    omni: AsyncOmni,
+async def run_inference_async_omni_llm(
+    omni_llm: AsyncOmniLLM,
     sampling_params: list[SamplingParams],
     prompt_args_template: SimpleNamespace,
     user_prompt: str,
@@ -246,7 +246,7 @@ async def run_inference_async_omni(
     video_file: Optional[str] = None,
     use_audio_in_video: bool = False,
 ):
-    """Run inference using AsyncOmni directly with multimodal support."""
+    """Run inference using AsyncOmniLLM directly with multimodal support."""
     if not user_prompt.strip() and not audio_file and not image_file and not video_file:
         return "Please provide at least a text prompt or multimodal input.", None
 
@@ -328,7 +328,7 @@ async def run_inference_async_omni(
         text_outputs: list[str] = []
         audio_output = None
 
-        async for stage_outputs in omni.generate(
+        async for stage_outputs in omni_llm.generate(
             prompt=omni_prompt,
             request_id=request_id,
             sampling_params_list=sampling_params,
@@ -362,12 +362,12 @@ async def run_inference_async_omni(
 
 
 def build_interface(
-    omni: AsyncOmni,
+    omni_llm: AsyncOmniLLM,
     sampling_params: list[SamplingParams],
     prompt_args_template: SimpleNamespace,
     model: str,
 ):
-    """Build Gradio interface for AsyncOmni mode."""
+    """Build Gradio interface for AsyncOmniLLM mode."""
 
     async def run_inference(
         user_prompt: str,
@@ -376,8 +376,8 @@ def build_interface(
         video_file: Optional[str],
         use_audio_in_video: bool,
     ):
-        return await run_inference_async_omni(
-            omni,
+        return await run_inference_async_omni_llm(
+            omni_llm,
             sampling_params,
             prompt_args_template,
             user_prompt,
@@ -475,7 +475,7 @@ def build_interface(
 
 def main():
     args = parse_args()
-    omni = None
+    omni_llm = None
 
     model_name = "/".join(args.model.split("/")[-2:])
     assert model_name in SUPPORTED_MODELS, (
@@ -485,9 +485,9 @@ def main():
     # Register signal handlers for graceful shutdown
     def signal_handler(sig, frame):
         print("\nReceived interrupt signal, shutting down...")
-        if omni is not None:
+        if omni_llm is not None:
             try:
-                omni.shutdown()
+                omni_llm.shutdown()
             except Exception as e:
                 print(f"Error during shutdown: {e}")
         sys.exit(0)
@@ -495,21 +495,21 @@ def main():
     signal.signal(signal.SIGINT, signal_handler)
     signal.signal(signal.SIGTERM, signal_handler)
 
-    print(f"Initializing AsyncOmni with model: {args.model}")
+    print(f"Initializing AsyncOmniLLM with model: {args.model}")
     if args.stage_configs_path:
         print(f"Using custom stage configs: {args.stage_configs_path}")
 
     sampling_params = build_sampling_params(SEED, model_name)
-    omni = AsyncOmni(
+    omni_llm = AsyncOmniLLM(
         model=args.model,
         stage_configs_path=args.stage_configs_path,
         init_timeout=ASYNC_INIT_TIMEOUT,
     )
-    print("✓ AsyncOmni initialized successfully")
+    print("✓ AsyncOmniLLM initialized successfully")
     prompt_args_template = create_prompt_args(args)
 
     demo = build_interface(
-        omni,
+        omni_llm,
         sampling_params,
         prompt_args_template,
         args.model,
@@ -524,9 +524,9 @@ def main():
         print("\nShutting down...")
     finally:
         # Cleanup
-        if omni is not None:
+        if omni_llm is not None:
             try:
-                omni.shutdown()
+                omni_llm.shutdown()
             except Exception as e:
                 print(f"Error during cleanup: {e}")
 

examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -1,6 +1,4 @@
 import base64
-import os
-from typing import Optional
 
 import requests
 from openai import OpenAI
@@ -30,129 +28,6 @@ def encode_base64_content_from_url(content_url: str) -> str:
     return result
 
 
-def encode_base64_content_from_file(file_path: str) -> str:
-    """Encode a local file to base64 format."""
-    with open(file_path, "rb") as f:
-        content = f.read()
-        result = base64.b64encode(content).decode("utf-8")
-    return result
-
-
-def get_video_url_from_path(video_path: Optional[str]) -> str:
-    """Convert a video path (local file or URL) to a video URL format for the API.
-
-    If video_path is None or empty, returns the default URL.
-    If video_path is a local file path, encodes it to base64 data URL.
-    If video_path is a URL, returns it as-is.
-    """
-    if not video_path:
-        # Default video URL
-        return "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-
-    # Check if it's a URL (starts with http:// or https://)
-    if video_path.startswith(("http://", "https://")):
-        return video_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(video_path):
-        raise FileNotFoundError(f"Video file not found: {video_path}")
-
-    # Detect video MIME type from file extension
-    video_path_lower = video_path.lower()
-    if video_path_lower.endswith(".mp4"):
-        mime_type = "video/mp4"
-    elif video_path_lower.endswith(".webm"):
-        mime_type = "video/webm"
-    elif video_path_lower.endswith(".mov"):
-        mime_type = "video/quicktime"
-    elif video_path_lower.endswith(".avi"):
-        mime_type = "video/x-msvideo"
-    elif video_path_lower.endswith(".mkv"):
-        mime_type = "video/x-matroska"
-    else:
-        # Default to mp4 if extension is unknown
-        mime_type = "video/mp4"
-
-    video_base64 = encode_base64_content_from_file(video_path)
-    return f"data:{mime_type};base64,{video_base64}"
-
-
-def get_image_url_from_path(image_path: Optional[str]) -> str:
-    """Convert an image path (local file or URL) to an image URL format for the API.
-
-    If image_path is None or empty, returns the default URL.
-    If image_path is a local file path, encodes it to base64 data URL.
-    If image_path is a URL, returns it as-is.
-    """
-    if not image_path:
-        # Default image URL
-        return "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-
-    # Check if it's a URL (starts with http:// or https://)
-    if image_path.startswith(("http://", "https://")):
-        return image_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(image_path):
-        raise FileNotFoundError(f"Image file not found: {image_path}")
-
-    # Detect image MIME type from file extension
-    image_path_lower = image_path.lower()
-    if image_path_lower.endswith((".jpg", ".jpeg")):
-        mime_type = "image/jpeg"
-    elif image_path_lower.endswith(".png"):
-        mime_type = "image/png"
-    elif image_path_lower.endswith(".gif"):
-        mime_type = "image/gif"
-    elif image_path_lower.endswith(".webp"):
-        mime_type = "image/webp"
-    else:
-        # Default to jpeg if extension is unknown
-        mime_type = "image/jpeg"
-
-    image_base64 = encode_base64_content_from_file(image_path)
-    return f"data:{mime_type};base64,{image_base64}"
-
-
-def get_audio_url_from_path(audio_path: Optional[str]) -> str:
-    """Convert an audio path (local file or URL) to an audio URL format for the API.
-
-    If audio_path is None or empty, returns the default URL.
-    If audio_path is a local file path, encodes it to base64 data URL.
-    If audio_path is a URL, returns it as-is.
-    """
-    if not audio_path:
-        # Default audio URL
-        return AudioAsset("mary_had_lamb").url
-
-    # Check if it's a URL (starts with http:// or https://)
-    if audio_path.startswith(("http://", "https://")):
-        return audio_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(audio_path):
-        raise FileNotFoundError(f"Audio file not found: {audio_path}")
-
-    # Detect audio MIME type from file extension
-    audio_path_lower = audio_path.lower()
-    if audio_path_lower.endswith((".mp3", ".mpeg")):
-        mime_type = "audio/mpeg"
-    elif audio_path_lower.endswith(".wav"):
-        mime_type = "audio/wav"
-    elif audio_path_lower.endswith(".ogg"):
-        mime_type = "audio/ogg"
-    elif audio_path_lower.endswith(".flac"):
-        mime_type = "audio/flac"
-    elif audio_path_lower.endswith(".m4a"):
-        mime_type = "audio/mp4"
-    else:
-        # Default to wav if extension is unknown
-        mime_type = "audio/wav"
-
-    audio_base64 = encode_base64_content_from_file(audio_path)
-    return f"data:{mime_type};base64,{audio_base64}"
-
-
 def get_system_prompt():
     return {
         "role": "system",
@@ -169,10 +44,8 @@ def get_system_prompt():
     }
 
 
-def get_text_query(custom_prompt: Optional[str] = None):
-    question = (
-        custom_prompt or "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
-    )
+def get_text_query():
+    question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
     prompt = {
         "role": "user",
         "content": [
@@ -185,32 +58,26 @@ def get_text_query(custom_prompt: Optional[str] = None):
     return prompt
 
 
-def get_mixed_modalities_query(
-    video_path: Optional[str] = None,
-    image_path: Optional[str] = None,
-    audio_path: Optional[str] = None,
-    custom_prompt: Optional[str] = None,
-):
-    question = (
-        custom_prompt or "What is recited in the audio? What is the content of this image? Why is this video funny?"
-    )
-    video_url = get_video_url_from_path(video_path)
-    image_url = get_image_url_from_path(image_path)
-    audio_url = get_audio_url_from_path(audio_path)
+def get_mixed_modalities_query():
+    question = "What is recited in the audio? What is the content of this image? Why is this video funny?"
     prompt = {
         "role": "user",
         "content": [
             {
                 "type": "audio_url",
-                "audio_url": {"url": audio_url},
+                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
             },
             {
                 "type": "image_url",
-                "image_url": {"url": image_url},
+                "image_url": {
+                    "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+                },
             },
             {
                 "type": "video_url",
-                "video_url": {"url": video_url},
+                "video_url": {
+                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+                },
             },
             {
                 "type": "text",
@@ -222,9 +89,8 @@ def get_mixed_modalities_query(
     return prompt
 
 
-def get_use_audio_in_video_query(video_path: Optional[str] = None, custom_prompt: Optional[str] = None):
-    question = custom_prompt or "Describe the content of the video, then convert what the baby say into text."
-    video_url = get_video_url_from_path(video_path)
+def get_use_audio_in_video_query():
+    question = "Describe the content of the video, then convert what the baby say into text."
 
     prompt = {
         "role": "user",
@@ -232,7 +98,7 @@ def get_use_audio_in_video_query(video_path: Optional[str] = None, custom_prompt
             {
                 "type": "video_url",
                 "video_url": {
-                    "url": video_url,
+                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4",
                     "num_frames": 16,
                 },
             },
@@ -246,15 +112,14 @@ def get_use_audio_in_video_query(video_path: Optional[str] = None, custom_prompt
     return prompt
 
 
-def get_multi_audios_query(audio_path: Optional[str] = None, custom_prompt: Optional[str] = None):
-    question = custom_prompt or "Are these two audio clips the same?"
-    audio_url = get_audio_url_from_path(audio_path)
+def get_multi_audios_query():
+    question = "Are these two audio clips the same?"
     prompt = {
         "role": "user",
         "content": [
             {
                 "type": "audio_url",
-                "audio_url": {"url": audio_url},
+                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
             },
             {
                 "type": "audio_url",
@@ -314,27 +179,7 @@ def run_multimodal_generation(args) -> None:
         code2wav_sampling_params,
     ]
 
-    # Get paths and custom prompt from args
-    video_path = getattr(args, "video_path", None)
-    image_path = getattr(args, "image_path", None)
-    audio_path = getattr(args, "audio_path", None)
-    custom_prompt = getattr(args, "prompt", None)
-
-    # Get the query function and call it with appropriate parameters
-    query_func = query_map[args.query_type]
-    if args.query_type == "mixed_modalities":
-        prompt = query_func(
-            video_path=video_path, image_path=image_path, audio_path=audio_path, custom_prompt=custom_prompt
-        )
-    elif args.query_type == "use_audio_in_video":
-        prompt = query_func(video_path=video_path, custom_prompt=custom_prompt)
-    elif args.query_type == "multi_audios":
-        prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
-    elif args.query_type == "text":
-        prompt = query_func(custom_prompt=custom_prompt)
-    else:
-        prompt = query_func()
-
+    prompt = query_map[args.query_type]()
     extra_body = {
         "sampling_params_list": sampling_params_list  # Optional, it has a default setting in stage_configs of the corresponding model.
     }
@@ -374,34 +219,6 @@ def parse_args():
         choices=query_map.keys(),
         help="Query type.",
     )
-    parser.add_argument(
-        "--video-path",
-        "-v",
-        type=str,
-        default=None,
-        help="Path to local video file or URL. If not provided and query-type uses video, uses default video URL.",
-    )
-    parser.add_argument(
-        "--image-path",
-        "-i",
-        type=str,
-        default=None,
-        help="Path to local image file or URL. If not provided and query-type uses image, uses default image URL.",
-    )
-    parser.add_argument(
-        "--audio-path",
-        "-a",
-        type=str,
-        default=None,
-        help="Path to local audio file or URL. If not provided and query-type uses audio, uses default audio URL.",
-    )
-    parser.add_argument(
-        "--prompt",
-        "-p",
-        type=str,
-        default=None,
-        help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
-    )
 
     return parser.parse_args()
 

examples/online_serving/qwen3_omni/README.md

@@ -1,13 +1,12 @@
-# Online serving Example of vLLM-Omni for Qwen3-omni
+# Online serving Example of vLLM-omni for Qwen2.5-omni
 
 ## 🛠️ Installation
 
-Please refer to [README.md](../../../README.md)
+Please refer to [README.md](../../README.md)
 
 ## Run examples (Qwen3-Omni)
 
-### Launch the Server
-
+Launch the server
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
@@ -17,59 +16,21 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 
-### Send Multi-modal Request
-
 Get into the example folder
 ```bash
-cd examples/online_serving/qwen3_omni
+cd examples/online_serving
 ```
 
-####  Send request via python
-
+Send request via python
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
+python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```
 
-The Python client supports the following command-line arguments:
-
-- `--query-type` (or `-q`): Query type (default: `use_video`)
-  - Options: `text`, `use_audio`, `use_image`, `use_video`
-- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-Omni-30B-A3B-Instruct`)
-- `--video-path` (or `-v`): Path to local video file or URL
-  - If not provided and query-type is `use_video`, uses default video URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
-- `--image-path` (or `-i`): Path to local image file or URL
-  - If not provided and query-type is `use_image`, uses default image URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common image formats: JPEG, PNG, GIF, WebP
-  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
-- `--audio-path` (or `-a`): Path to local audio file or URL
-  - If not provided and query-type is `use_audio`, uses default audio URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
-  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
-- `--prompt` (or `-p`): Custom text prompt/question
-  - If not provided, uses default prompt for the selected query type
-  - Example: `--prompt "What are the main activities shown in this video?"`
-
-
-For example, to use a local video file with custom prompt:
-
+Send request via curl
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py \
-    --query-type use_video \
-    --video-path /path/to/your/video.mp4 \
-    --prompt "What are the main activities shown in this video?"
+bash run_curl_multimodal_generation.sh mixed_modalities
 ```
 
-####  Send request via curl
-
-```bash
-bash run_curl_multimodal_generation.sh use_image
-```
-
-
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -84,7 +45,7 @@ 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 AsyncOmni by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
@@ -93,9 +54,23 @@ python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
 Then open `http://localhost:7861/` on your local browser to interact with the web UI.
 
 
-The gradio script supports the following arguments:
+### 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/Qwen3-Omni-30B-A3B-Instruct \
+    --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)

examples/online_serving/qwen3_omni/gradio_demo.py

@@ -15,7 +15,7 @@ from PIL import Image
 from vllm.assets.video import video_get_metadata, video_to_ndarrays
 from vllm.sampling_params import SamplingParams
 
-from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
 
 # Import utils from offline inference example
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen3_omni"))
@@ -42,20 +42,19 @@ SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
                 "seed": SEED,
             },
             "talker": {
-                "temperature": 0.9,
-                "top_k": 50,
+                "temperature": 0.0,
+                "top_p": 1.0,
+                "top_k": -1,
                 "max_tokens": 4096,
-                "seed": SEED,
-                "detokenize": False,
-                "repetition_penalty": 1.05,
+                "detokenize": True,
+                "repetition_penalty": 1.1,
                 "stop_token_ids": [2150],
             },
             "code2wav": {
                 "temperature": 0.0,
                 "top_p": 1.0,
                 "top_k": -1,
-                "max_tokens": 4096 * 16,
-                "seed": SEED,
+                "max_tokens": 2048,
                 "detokenize": True,
                 "repetition_penalty": 1.1,
             },
@@ -239,8 +238,8 @@ def process_video_file(
     return frames, metadata, audio_tuple
 
 
-async def run_inference_async_omni(
-    omni: AsyncOmni,
+async def run_inference_async_omni_llm(
+    omni_llm: AsyncOmniLLM,
     sampling_params: list[SamplingParams],
     prompt_args_template: SimpleNamespace,
     user_prompt: str,
@@ -249,7 +248,7 @@ async def run_inference_async_omni(
     video_file: Optional[str] = None,
     use_audio_in_video: bool = False,
 ):
-    """Run inference using AsyncOmni directly with multimodal support."""
+    """Run inference using AsyncOmniLLM directly with multimodal support."""
     if not user_prompt.strip() and not audio_file and not image_file and not video_file:
         return "Please provide at least a text prompt or multimodal input.", None
 
@@ -264,7 +263,7 @@ async def run_inference_async_omni(
         # Process audio
         audio_data = process_audio_file(audio_file)
         if audio_data is not None:
-            multi_modal_data["audio"] = [audio_data] if not isinstance(audio_data, list) else audio_data
+            multi_modal_data["audio"] = audio_data
 
         # Process image
         image_data = process_image_file(image_file)
@@ -282,9 +281,7 @@ async def run_inference_async_omni(
                 video_entry = video_frames
             multi_modal_data["video"] = video_entry
             if use_audio_in_video and extracted_audio is not None and "audio" not in multi_modal_data:
-                multi_modal_data["audio"] = (
-                    [extracted_audio] if not isinstance(extracted_audio, list) else extracted_audio
-                )
+                multi_modal_data["audio"] = extracted_audio
                 mm_processor_kwargs["use_audio_in_video"] = True
 
         # Build the prompt input
@@ -313,11 +310,11 @@ async def run_inference_async_omni(
             )
             prompt = f"<|im_start|>system\n{default_system}<|im_end|>\n<|im_start|>user\n"
             if audio_data:
-                prompt += "<|audio_start|><|audio_pad|><|audio_end|>"
+                prompt += "<|audio_bos|><|AUDIO|><|audio_eos|>"
             if image_data:
-                prompt += "<|vision_start|><|image_pad|><|vision_end|>"
+                prompt += "<|vision_bos|><|IMAGE|><|vision_eos|>"
             if video_payload is not None:
-                prompt += "<|vision_start|><|video_pad|><|vision_end|>"
+                prompt += "<|vision_bos|><|VIDEO|><|vision_eos|>"
             if user_prompt.strip():
                 prompt += f"{user_prompt}"
             prompt += "<|im_end|>\n<|im_start|>assistant\n"
@@ -333,7 +330,7 @@ async def run_inference_async_omni(
         text_outputs: list[str] = []
         audio_output = None
 
-        async for stage_outputs in omni.generate(
+        async for stage_outputs in omni_llm.generate(
             prompt=omni_prompt,
             request_id=request_id,
             sampling_params_list=sampling_params,
@@ -347,15 +344,14 @@ async def run_inference_async_omni(
                             text_outputs.append(output.text)
             elif stage_outputs.final_output_type == "audio":
                 # multimodal_output is on the RequestOutput object
-                # See vllm_omni/entrypoints/openai/serving_chat.py:733 for reference
+                # See vllm_omni/entrypoints/openai/serving_chat.py:680 for reference
                 if hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
                     if "audio" in request_output.multimodal_output:
                         audio_tensor = request_output.multimodal_output["audio"]
                         # Ensure audio is 1D (flatten if needed)
                         if hasattr(audio_tensor, "ndim") and audio_tensor.ndim > 1:
                             audio_tensor = audio_tensor.flatten()
-                        # Convert BFloat16 to Float32 before numpy conversion (numpy doesn't support BFloat16)
-                        audio_np = audio_tensor.float().detach().cpu().numpy()
+                        audio_np = audio_tensor.detach().cpu().numpy()
                         audio_output = (
                             24000,  # sampling rate in Hz
                             audio_np,
@@ -368,12 +364,12 @@ async def run_inference_async_omni(
 
 
 def build_interface(
-    omni: AsyncOmni,
+    omni_llm: AsyncOmniLLM,
     sampling_params: list[SamplingParams],
     prompt_args_template: SimpleNamespace,
     model: str,
 ):
-    """Build Gradio interface for AsyncOmni mode."""
+    """Build Gradio interface for AsyncOmniLLM mode."""
 
     async def run_inference(
         user_prompt: str,
@@ -382,8 +378,8 @@ def build_interface(
         video_file: Optional[str],
         use_audio_in_video: bool,
     ):
-        return await run_inference_async_omni(
-            omni,
+        return await run_inference_async_omni_llm(
+            omni_llm,
             sampling_params,
             prompt_args_template,
             user_prompt,
@@ -481,7 +477,7 @@ def build_interface(
 
 def main():
     args = parse_args()
-    omni = None
+    omni_llm = None
 
     model_name = "/".join(args.model.split("/")[-2:])
     assert model_name in SUPPORTED_MODELS, (
@@ -491,9 +487,9 @@ def main():
     # Register signal handlers for graceful shutdown
     def signal_handler(sig, frame):
         print("\nReceived interrupt signal, shutting down...")
-        if omni is not None:
+        if omni_llm is not None:
             try:
-                omni.shutdown()
+                omni_llm.shutdown()
             except Exception as e:
                 print(f"Error during shutdown: {e}")
         sys.exit(0)
@@ -501,21 +497,21 @@ def main():
     signal.signal(signal.SIGINT, signal_handler)
     signal.signal(signal.SIGTERM, signal_handler)
 
-    print(f"Initializing AsyncOmni with model: {args.model}")
+    print(f"Initializing AsyncOmniLLM with model: {args.model}")
     if args.stage_configs_path:
         print(f"Using custom stage configs: {args.stage_configs_path}")
 
     sampling_params = build_sampling_params(SEED, model_name)
-    omni = AsyncOmni(
+    omni_llm = AsyncOmniLLM(
         model=args.model,
         stage_configs_path=args.stage_configs_path,
         init_timeout=ASYNC_INIT_TIMEOUT,
     )
-    print("✓ AsyncOmni initialized successfully")
+    print("✓ AsyncOmniLLM initialized successfully")
     prompt_args_template = create_prompt_args(args)
 
     demo = build_interface(
-        omni,
+        omni_llm,
         sampling_params,
         prompt_args_template,
         args.model,
@@ -530,9 +526,9 @@ def main():
         print("\nShutting down...")
     finally:
         # Cleanup
-        if omni is not None:
+        if omni_llm is not None:
             try:
-                omni.shutdown()
+                omni_llm.shutdown()
             except Exception as e:
                 print(f"Error during cleanup: {e}")
 

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -1,6 +1,5 @@
 import base64
-import os
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 
 import requests
 from openai import OpenAI
@@ -35,129 +34,6 @@ def encode_base64_content_from_url(content_url: str) -> str:
     return result
 
 
-def encode_base64_content_from_file(file_path: str) -> str:
-    """Encode a local file to base64 format."""
-    with open(file_path, "rb") as f:
-        content = f.read()
-        result = base64.b64encode(content).decode("utf-8")
-    return result
-
-
-def get_video_url_from_path(video_path: Optional[str]) -> str:
-    """Convert a video path (local file or URL) to a video URL format for the API.
-
-    If video_path is None or empty, returns the default URL.
-    If video_path is a local file path, encodes it to base64 data URL.
-    If video_path is a URL, returns it as-is.
-    """
-    if not video_path:
-        # Default video URL
-        return "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-
-    # Check if it's a URL (starts with http:// or https://)
-    if video_path.startswith(("http://", "https://")):
-        return video_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(video_path):
-        raise FileNotFoundError(f"Video file not found: {video_path}")
-
-    # Detect video MIME type from file extension
-    video_path_lower = video_path.lower()
-    if video_path_lower.endswith(".mp4"):
-        mime_type = "video/mp4"
-    elif video_path_lower.endswith(".webm"):
-        mime_type = "video/webm"
-    elif video_path_lower.endswith(".mov"):
-        mime_type = "video/quicktime"
-    elif video_path_lower.endswith(".avi"):
-        mime_type = "video/x-msvideo"
-    elif video_path_lower.endswith(".mkv"):
-        mime_type = "video/x-matroska"
-    else:
-        # Default to mp4 if extension is unknown
-        mime_type = "video/mp4"
-
-    video_base64 = encode_base64_content_from_file(video_path)
-    return f"data:{mime_type};base64,{video_base64}"
-
-
-def get_image_url_from_path(image_path: Optional[str]) -> str:
-    """Convert an image path (local file or URL) to an image URL format for the API.
-
-    If image_path is None or empty, returns the default URL.
-    If image_path is a local file path, encodes it to base64 data URL.
-    If image_path is a URL, returns it as-is.
-    """
-    if not image_path:
-        # Default image URL
-        return "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-
-    # Check if it's a URL (starts with http:// or https://)
-    if image_path.startswith(("http://", "https://")):
-        return image_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(image_path):
-        raise FileNotFoundError(f"Image file not found: {image_path}")
-
-    # Detect image MIME type from file extension
-    image_path_lower = image_path.lower()
-    if image_path_lower.endswith((".jpg", ".jpeg")):
-        mime_type = "image/jpeg"
-    elif image_path_lower.endswith(".png"):
-        mime_type = "image/png"
-    elif image_path_lower.endswith(".gif"):
-        mime_type = "image/gif"
-    elif image_path_lower.endswith(".webp"):
-        mime_type = "image/webp"
-    else:
-        # Default to jpeg if extension is unknown
-        mime_type = "image/jpeg"
-
-    image_base64 = encode_base64_content_from_file(image_path)
-    return f"data:{mime_type};base64,{image_base64}"
-
-
-def get_audio_url_from_path(audio_path: Optional[str]) -> str:
-    """Convert an audio path (local file or URL) to an audio URL format for the API.
-
-    If audio_path is None or empty, returns the default URL.
-    If audio_path is a local file path, encodes it to base64 data URL.
-    If audio_path is a URL, returns it as-is.
-    """
-    if not audio_path:
-        # Default audio URL
-        return AudioAsset("mary_had_lamb").url
-
-    # Check if it's a URL (starts with http:// or https://)
-    if audio_path.startswith(("http://", "https://")):
-        return audio_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(audio_path):
-        raise FileNotFoundError(f"Audio file not found: {audio_path}")
-
-    # Detect audio MIME type from file extension
-    audio_path_lower = audio_path.lower()
-    if audio_path_lower.endswith((".mp3", ".mpeg")):
-        mime_type = "audio/mpeg"
-    elif audio_path_lower.endswith(".wav"):
-        mime_type = "audio/wav"
-    elif audio_path_lower.endswith(".ogg"):
-        mime_type = "audio/ogg"
-    elif audio_path_lower.endswith(".flac"):
-        mime_type = "audio/flac"
-    elif audio_path_lower.endswith(".m4a"):
-        mime_type = "audio/mp4"
-    else:
-        # Default to wav if extension is unknown
-        mime_type = "audio/wav"
-
-    audio_base64 = encode_base64_content_from_file(audio_path)
-    return f"data:{mime_type};base64,{audio_base64}"
-
-
 def get_system_prompt():
     return {
         "role": "system",
@@ -174,10 +50,8 @@ def get_system_prompt():
     }
 
 
-def get_text_query(custom_prompt: Optional[str] = None):
-    question = (
-        custom_prompt or "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
-    )
+def get_text_query():
+    question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
     prompt = {
         "role": "user",
         "content": [
@@ -197,15 +71,16 @@ default_system = (
 )
 
 
-def get_video_query(video_path: Optional[str] = None, custom_prompt: Optional[str] = None):
-    question = custom_prompt or "Why is this video funny?"
-    video_url = get_video_url_from_path(video_path)
+def get_video_query() -> QueryResult:
+    question = "Why is this video funny?"
     prompt = {
         "role": "user",
         "content": [
             {
                 "type": "video_url",
-                "video_url": {"url": video_url},
+                "video_url": {
+                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+                },
             },
             {
                 "type": "text",
@@ -216,15 +91,16 @@ def get_video_query(video_path: Optional[str] = None, custom_prompt: Optional[st
     return prompt
 
 
-def get_image_query(image_path: Optional[str] = None, custom_prompt: Optional[str] = None):
-    question = custom_prompt or "What is the content of this image?"
-    image_url = get_image_url_from_path(image_path)
+def get_image_query() -> QueryResult:
+    question = "What is the content of this image?"
     prompt = {
         "role": "user",
         "content": [
             {
                 "type": "image_url",
-                "image_url": {"url": image_url},
+                "image_url": {
+                    "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+                },
             },
             {
                 "type": "text",
@@ -235,15 +111,14 @@ def get_image_query(image_path: Optional[str] = None, custom_prompt: Optional[st
     return prompt
 
 
-def get_audio_query(audio_path: Optional[str] = None, custom_prompt: Optional[str] = None):
-    question = custom_prompt or "What is the content of this audio?"
-    audio_url = get_audio_url_from_path(audio_path)
+def get_audio_query() -> QueryResult:
+    question = "What is the content of this audio?"
     prompt = {
         "role": "user",
         "content": [
             {
                 "type": "audio_url",
-                "audio_url": {"url": audio_url},
+                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
             },
             {
                 "type": "text",
@@ -303,25 +178,7 @@ def run_multimodal_generation(args) -> None:
         code2wav_sampling_params,
     ]
 
-    # Get paths and custom prompt from args
-    video_path = getattr(args, "video_path", None)
-    image_path = getattr(args, "image_path", None)
-    audio_path = getattr(args, "audio_path", None)
-    custom_prompt = getattr(args, "prompt", None)
-
-    # Get the query function and call it with appropriate parameters
-    query_func = query_map[args.query_type]
-    if args.query_type == "use_video":
-        prompt = query_func(video_path=video_path, custom_prompt=custom_prompt)
-    elif args.query_type == "use_image":
-        prompt = query_func(image_path=image_path, custom_prompt=custom_prompt)
-    elif args.query_type == "use_audio":
-        prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
-    elif args.query_type == "text":
-        prompt = query_func(custom_prompt=custom_prompt)
-    else:
-        prompt = query_func()
-
+    prompt = query_map[args.query_type]()
     extra_body = {
         "sampling_params_list": sampling_params_list  # Optional, it has a default setting in stage_configs of the corresponding model.
     }
@@ -340,6 +197,7 @@ def run_multimodal_generation(args) -> None:
 
     count = 0
     for choice in chat_completion.choices:
+        print(f"choice: {choice}")
         if choice.message.audio:
             audio_data = base64.b64decode(choice.message.audio.data)
             audio_file_path = f"audio_{count}.wav"
@@ -368,34 +226,6 @@ def parse_args():
         default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
         help="Model Name / Path",
     )
-    parser.add_argument(
-        "--video-path",
-        "-v",
-        type=str,
-        default=None,
-        help="Path to local video file or URL. If not provided and query-type is 'use_video', uses default video URL.",
-    )
-    parser.add_argument(
-        "--image-path",
-        "-i",
-        type=str,
-        default=None,
-        help="Path to local image file or URL. If not provided and query-type is 'use_image', uses default image URL.",
-    )
-    parser.add_argument(
-        "--audio-path",
-        "-a",
-        type=str,
-        default=None,
-        help="Path to local audio file or URL. If not provided and query-type is 'use_audio', uses default audio URL.",
-    )
-    parser.add_argument(
-        "--prompt",
-        "-p",
-        type=str,
-        default=None,
-        help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
-    )
 
     return parser.parse_args()
 

examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml

@@ -1,36 +0,0 @@
-# Stage config for running Qwen3-Omni-MoE-Thinking (text-only output)
-# This config is for models like Qwen3-Omni-30B-A3B-Thinking that only have the
-# thinker component and do not support audio output.
-#
-# Single stage: Thinker (multimodal understanding + text generation)
-
-# The following config has been verified on 2x H100-80G GPUs.
-stage_args:
-  - stage_id: 0
-    runtime:
-      devices: "0,1"
-      max_batch_size: 1
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: text
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      hf_config_name: thinker_config
-      tensor_parallel_size: 2
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05

examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh

@@ -2,16 +2,16 @@
 set -euo pipefail
 
 # Default query type
-QUERY_TYPE="${1:-use_video}"
+QUERY_TYPE="${1:-mixed_modalities}"
 
 # Validate query type
-if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video)$ ]]; then
+if [[ ! "$QUERY_TYPE" =~ ^(mixed_modalities|use_audio_in_video|multi_audios|text)$ ]]; then
     echo "Error: Invalid query type '$QUERY_TYPE'"
-    echo "Usage: $0 [text|use_audio|use_image|use_video]"
+    echo "Usage: $0 [mixed_modalities|use_audio_in_video|multi_audios|text]"
+    echo "  mixed_modalities: Audio + Image + Video + Text query"
+    echo "  use_audio_in_video: Video + Text query (with audio extraction from video)"
+    echo "  multi_audios: Two audio clips + Text query"
     echo "  text: Text query"
-    echo "  use_audio: Audio + Text query"
-    echo "  use_image: Image + Text query"
-    echo "  use_video: Video + Text query"
     exit 1
 fi
 
@@ -23,6 +23,7 @@ thinker_sampling_params='{
   "top_k": 1,
   "max_tokens": 16384,
   "seed": 42,
+  "detokenize": true,
   "repetition_penalty": 1.05,
   "stop_token_ids": [151645]
 }'
@@ -32,7 +33,7 @@ talker_sampling_params='{
   "top_k": 50,
   "max_tokens": 4096,
   "seed": 42,
-  "detokenize": false,
+  "detokenize": true,
   "repetition_penalty": 1.05,
   "stop_token_ids": [2150]
 }'
@@ -41,7 +42,7 @@ code2wav_sampling_params='{
   "temperature": 0.0,
   "top_p": 1.0,
   "top_k": -1,
-  "max_tokens": 65536,
+  "max_tokens": 4096 * 16,
   "seed": 42,
   "detokenize": true,
   "repetition_penalty": 1.1
@@ -50,6 +51,7 @@ code2wav_sampling_params='{
 
 # Define URLs for assets
 MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
+WINNING_CALL_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/winning_call.ogg"
 CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
 SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
 
@@ -69,7 +71,7 @@ case "$QUERY_TYPE" in
     ]'
     mm_processor_kwargs="{}"
     ;;
-  use_audio)
+  mixed_modalities)
     user_content='[
         {
           "type": "audio_url",
@@ -77,29 +79,21 @@ case "$QUERY_TYPE" in
             "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
           }
         },
-        {
-          "type": "text",
-          "text": "What is the content of this audio?"
-        }
-      ]'
-    sampling_params_list='[
-      '"$thinker_sampling_params"',
-      '"$talker_sampling_params"',
-      '"$code2wav_sampling_params"'
-    ]'
-    mm_processor_kwargs="{}"
-    ;;
-  use_image)
-    user_content='[
         {
           "type": "image_url",
           "image_url": {
             "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
           }
         },
+        {
+          "type": "video_url",
+          "video_url": {
+            "url": "'"$SAMPLE_VIDEO_URL"'"
+          }
+        },
         {
           "type": "text",
-          "text": "What is the content of this image?"
+          "text": "What is recited in the audio? What is the content of this image? Why is this video funny?"
         }
       ]'
     sampling_params_list='[
@@ -109,7 +103,7 @@ case "$QUERY_TYPE" in
     ]'
     mm_processor_kwargs="{}"
     ;;
-  use_video)
+  use_audio_in_video)
     user_content='[
         {
           "type": "video_url",
@@ -119,7 +113,35 @@ case "$QUERY_TYPE" in
         },
         {
           "type": "text",
-          "text": "Why is this video funny?"
+          "text": "Describe the content of the video, then convert what the baby say into text."
+        }
+      ]'
+    sampling_params_list='[
+      '"$thinker_sampling_params"',
+      '"$talker_sampling_params"',
+      '"$code2wav_sampling_params"'
+    ]'
+    mm_processor_kwargs='{
+      "use_audio_in_video": true
+    }'
+    ;;
+  multi_audios)
+    user_content='[
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
+          }
+        },
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "'"$WINNING_CALL_AUDIO_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Are these two audio clips the same?"
         }
       ]'
     sampling_params_list='[
@@ -139,7 +161,7 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
     -H "Content-Type: application/json" \
     -d @- <<EOF
 {
-  "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+  "model": "Qwen/Qwen2.5-Omni-7B",
   "sampling_params_list": $sampling_params_list,
   "mm_processor_kwargs": $mm_processor_kwargs,
   "messages": [

mkdocs.yml

@@ -1,5 +1,5 @@
 site_name: vLLM-Omni
-site_description: Efficient omni-modality model serving for everyone
+site_description: Multi-modality models inference and serving with non-autoregressive structures
 site_author: vLLM-Omni Team
 site_url: https://vllm-project.github.io/vllm-omni/
 
@@ -58,7 +58,6 @@ theme:
 hooks:
   - docs/mkdocs/hooks/generate_api_readme.py
   - docs/mkdocs/hooks/url_schemes.py
-  - docs/mkdocs/hooks/generate_examples.py
 
 # Plugins
 plugins:
@@ -88,7 +87,6 @@ plugins:
       show_full_namespace: false  # Show only module name, not full path
       exclude:
         - "re:vllm_omni\\._.*"  # Internal modules
-        - "vllm_omni.diffusion.models.qwen_image"  # avoid importing vllm in mkdocs building
   - mkdocstrings:
       handlers:
         python:

pyproject.toml

@@ -5,12 +5,12 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "vllm-omni"
 version = "0.11.0rc1"
-description = "A framework for efficient model inference with omni-modality models"
+description = "vLLM-Omni: Multi-modality models inference and serving with non-autoregressive structures"
 readme = "README.md"
 requires-python = ">=3.9,<3.14"  # Align with vLLM v0.11 supported Python versions
 license = {text = "Apache-2.0"}
 authors = [
-    {name = "vLLM-Omni Team"}
+    {name = "vLLM-Omni Team", email = "vllm-project@gmail.com"}
 ]
 keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
 classifiers = [
@@ -27,7 +27,6 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 
-
 dependencies = [
     # Core runtime dependencies (required for actual usage)
     "omegaconf>=2.3.0",
@@ -156,7 +155,6 @@ markers = [
     "integration: Integration tests",
     "benchmark: Benchmark tests",
     "slow: Slow tests",
-    "core_model: enable this model test in each PR instead of only nightly",
 ]
 
 [tool.typos.default]

tests/__init__.py

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

tests/conftest.py

@@ -1,36 +0,0 @@
-import os
-
-import pytest
-import torch
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-@pytest.fixture(autouse=True)
-def clean_gpu_memory_between_tests():
-    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
-        yield
-        return
-
-    # Wait for GPU memory to be cleared before starting the test
-    import gc
-
-    from tests.utils import wait_for_gpu_memory_to_clear
-
-    num_gpus = torch.cuda.device_count()
-    if num_gpus > 0:
-        try:
-            wait_for_gpu_memory_to_clear(
-                devices=list(range(num_gpus)),
-                threshold_ratio=0.1,
-            )
-        except ValueError as e:
-            logger.info("Failed to clean GPU memory: %s", e)
-
-    yield
-
-    # Clean up GPU memory after the test
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        gc.collect()

tests/multi_stages/conftest.py

@@ -1,341 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Pytest configuration and fixtures for vllm-omni tests.
-"""
-
-from typing import Any
-
-import pytest
-from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-from vllm.sampling_params import SamplingParams
-
-from vllm_omni.entrypoints.omni import Omni
-
-PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
-PromptImageInput = list[Any] | Any | None
-PromptVideoInput = list[Any] | Any | None
-
-
-class OmniRunner:
-    """
-    Test runner for Omni models.
-    """
-
-    def __init__(
-        self,
-        model_name: str,
-        seed: int = 42,
-        init_sleep_seconds: int = 20,
-        batch_timeout: int = 10,
-        init_timeout: int = 300,
-        shm_threshold_bytes: int = 65536,
-        log_stats: bool = False,
-        stage_configs_path: str | None = None,
-        **kwargs,
-    ) -> None:
-        """
-        Initialize an OmniRunner for testing.
-
-        Args:
-            model_name: The model name or path
-            seed: Random seed for reproducibility
-            init_sleep_seconds: Sleep time after starting each stage
-            batch_timeout: Timeout for batching in seconds
-            init_timeout: Timeout for initializing stages in seconds
-            shm_threshold_bytes: Threshold for using shared memory
-            log_stats: Enable detailed statistics logging
-            stage_configs_path: Optional path to YAML stage config file
-            **kwargs: Additional arguments passed to Omni
-        """
-        self.model_name = model_name
-        self.seed = seed
-
-        self.omni = Omni(
-            model=model_name,
-            log_stats=log_stats,
-            init_sleep_seconds=init_sleep_seconds,
-            batch_timeout=batch_timeout,
-            init_timeout=init_timeout,
-            shm_threshold_bytes=shm_threshold_bytes,
-            stage_configs_path=stage_configs_path,
-            **kwargs,
-        )
-
-    def get_default_sampling_params_list(self) -> list[SamplingParams]:
-        """
-        Get a list of default sampling parameters for all stages.
-
-        Returns:
-            List of SamplingParams with default decoding for each stage
-        """
-        return [st.default_sampling_params for st in self.omni.instance.stage_list]
-
-    def get_omni_inputs(
-        self,
-        prompts: list[str] | str,
-        system_prompt: str | None = None,
-        audios: PromptAudioInput = None,
-        images: PromptImageInput = None,
-        videos: PromptVideoInput = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> list[dict[str, Any]]:
-        """
-        Construct Omni input format from prompts and multimodal data.
-
-        Args:
-            prompts: Text prompt(s) - either a single string or list of strings
-            system_prompt: Optional system prompt (defaults to Qwen system prompt)
-            audios: Audio input(s) - tuple of (audio_array, sample_rate) or list of tuples
-            images: Image input(s) - PIL Image or list of PIL Images
-            videos: Video input(s) - numpy array or list of numpy arrays
-            mm_processor_kwargs: Optional processor kwargs (e.g., use_audio_in_video)
-
-        Returns:
-            List of prompt dictionaries suitable for Omni.generate()
-        """
-        if system_prompt is None:
-            system_prompt = (
-                "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
-                "Group, capable of perceiving auditory and visual inputs, as well as "
-                "generating text and speech."
-            )
-
-        video_padding_token = "<|VIDEO|>"
-        image_padding_token = "<|IMAGE|>"
-        audio_padding_token = "<|AUDIO|>"
-
-        if self.model_name == "Qwen/Qwen3-Omni-30B-A3B-Instruct":
-            video_padding_token = "<|video_pad|>"
-            image_padding_token = "<|image_pad|>"
-            audio_padding_token = "<|audio_pad|>"
-
-        if isinstance(prompts, str):
-            prompts = [prompts]
-
-        def _normalize_mm_input(mm_input, num_prompts):
-            if mm_input is None:
-                return [None] * num_prompts
-            if isinstance(mm_input, list):
-                if len(mm_input) != num_prompts:
-                    raise ValueError(
-                        f"Multimodal input list length ({len(mm_input)}) must match prompts length ({num_prompts})"
-                    )
-                return mm_input
-            return [mm_input] * num_prompts
-
-        num_prompts = len(prompts)
-        audios_list = _normalize_mm_input(audios, num_prompts)
-        images_list = _normalize_mm_input(images, num_prompts)
-        videos_list = _normalize_mm_input(videos, num_prompts)
-
-        omni_inputs = []
-        for i, prompt_text in enumerate(prompts):
-            user_content = ""
-            multi_modal_data = {}
-
-            audio = audios_list[i]
-            if audio is not None:
-                if isinstance(audio, list):
-                    for _ in audio:
-                        user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
-                    multi_modal_data["audio"] = audio
-                else:
-                    user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
-                    multi_modal_data["audio"] = audio
-
-            image = images_list[i]
-            if image is not None:
-                if isinstance(image, list):
-                    for _ in image:
-                        user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
-                    multi_modal_data["image"] = image
-                else:
-                    user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
-                    multi_modal_data["image"] = image
-
-            video = videos_list[i]
-            if video is not None:
-                if isinstance(video, list):
-                    for _ in video:
-                        user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
-                    multi_modal_data["video"] = video
-                else:
-                    user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
-                    multi_modal_data["video"] = video
-
-            user_content += prompt_text
-
-            full_prompt = (
-                f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
-                f"<|im_start|>user\n{user_content}<|im_end|>\n"
-                f"<|im_start|>assistant\n"
-            )
-
-            input_dict: dict[str, Any] = {"prompt": full_prompt}
-            if multi_modal_data:
-                input_dict["multi_modal_data"] = multi_modal_data
-            if mm_processor_kwargs:
-                input_dict["mm_processor_kwargs"] = mm_processor_kwargs
-
-            omni_inputs.append(input_dict)
-
-        return omni_inputs
-
-    def generate(
-        self,
-        prompts: list[dict[str, Any]],
-        sampling_params_list: list[SamplingParams] | None = None,
-    ) -> list[Any]:
-        """
-        Generate outputs for the given prompts.
-
-        Args:
-            prompts: List of prompt dictionaries with 'prompt' and optionally
-                    'multi_modal_data' keys
-            sampling_params_list: List of sampling parameters for each stage.
-                                 If None, uses default parameters.
-
-        Returns:
-            List of OmniRequestOutput objects from stages with final_output=True
-        """
-        if sampling_params_list is None:
-            sampling_params_list = self.get_default_sampling_params_list()
-
-        return self.omni.generate(prompts, sampling_params_list)
-
-    def generate_multimodal(
-        self,
-        prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
-        system_prompt: str | None = None,
-        audios: PromptAudioInput = None,
-        images: PromptImageInput = None,
-        videos: PromptVideoInput = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> list[Any]:
-        """
-        Convenience method to generate with multimodal inputs.
-
-        Args:
-            prompts: Text prompt(s)
-            sampling_params_list: List of sampling parameters for each stage
-            system_prompt: Optional system prompt
-            audios: Audio input(s)
-            images: Image input(s)
-            videos: Video input(s)
-            mm_processor_kwargs: Optional processor kwargs
-
-        Returns:
-            List of OmniRequestOutput objects from stages with final_output=True
-        """
-        omni_inputs = self.get_omni_inputs(
-            prompts=prompts,
-            system_prompt=system_prompt,
-            audios=audios,
-            images=images,
-            videos=videos,
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-        return self.generate(omni_inputs, sampling_params_list)
-
-    def generate_audio(
-        self,
-        prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
-        system_prompt: str | None = None,
-        audios: PromptAudioInput = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> list[Any]:
-        """
-        Convenience method to generate with multimodal inputs.
-        Args:
-            prompts: Text prompt(s)
-            sampling_params_list: List of sampling parameters for each stage
-            system_prompt: Optional system prompt
-            audios: Audio input(s)
-            mm_processor_kwargs: Optional processor kwargs
-        Returns:
-            List of OmniRequestOutput objects from stages with final_output=True
-        """
-        omni_inputs = self.get_omni_inputs(
-            prompts=prompts,
-            system_prompt=system_prompt,
-            audios=audios,
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-        return self.generate(omni_inputs, sampling_params_list)
-
-    def generate_video(
-        self,
-        prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
-        system_prompt: str | None = None,
-        videos: PromptVideoInput = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> list[Any]:
-        """
-        Convenience method to generate with multimodal inputs.
-        Args:
-            prompts: Text prompt(s)
-            sampling_params_list: List of sampling parameters for each stage
-            system_prompt: Optional system prompt
-            videos: Video input(s)
-            mm_processor_kwargs: Optional processor kwargs
-        Returns:
-            List of OmniRequestOutput objects from stages with final_output=True
-        """
-        omni_inputs = self.get_omni_inputs(
-            prompts=prompts,
-            system_prompt=system_prompt,
-            videos=videos,
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-        return self.generate(omni_inputs, sampling_params_list)
-
-    def generate_image(
-        self,
-        prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
-        system_prompt: str | None = None,
-        images: PromptImageInput = None,
-        mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> list[Any]:
-        """
-        Convenience method to generate with multimodal inputs.
-        Args:
-            prompts: Text prompt(s)
-            sampling_params_list: List of sampling parameters for each stage
-            system_prompt: Optional system prompt
-            images: Image input(s)
-            mm_processor_kwargs: Optional processor kwargs
-        Returns:
-            List of OmniRequestOutput objects from stages with final_output=True
-        """
-        omni_inputs = self.get_omni_inputs(
-            prompts=prompts,
-            system_prompt=system_prompt,
-            images=images,
-            mm_processor_kwargs=mm_processor_kwargs,
-        )
-        return self.generate(omni_inputs, sampling_params_list)
-
-    def __enter__(self):
-        """Context manager entry."""
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """Context manager exit - cleanup resources."""
-        self.close()
-        del self.omni
-        cleanup_dist_env_and_memory()
-
-    def close(self):
-        """Close and cleanup the Omni instance."""
-        if hasattr(self.omni.instance, "close"):
-            self.omni.instance.close()
-
-
-@pytest.fixture(scope="session")
-def omni_runner():
-    return OmniRunner

tests/multi_stages/stage_configs/qwen2_5_omni_ci.yaml

@@ -1,105 +0,0 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
-
-# The following config has been verified on 1x 24GB GPU (L4/RTX3090).
-# This config is optimized for CI e2e tests.
-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.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 896
-      max_num_batched_tokens: 896
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.8
-      skip_mm_profiling: true
-      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: 128
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1
-  - stage_id: 1
-    runtime:
-      process: true
-      devices: "1"
-      max_batch_size: 1
-    engine_args:
-      model_stage: talker
-      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 896
-      max_num_batched_tokens: 896
-      max_num_seqs: 1
-      gpu_memory_utilization: 0.8
-      skip_mm_profiling: true
-      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: 128
-      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_generation_worker.GPUGenerationWorker
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      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: 128
-      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

tests/multi_stages/test_qwen2_5_omni.py

@@ -1,73 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E tests for Qwen2.5-Omni model with mixed modality inputs and audio output.
-"""
-
-from pathlib import Path
-
-import pytest
-from vllm.assets.audio import AudioAsset
-from vllm.assets.image import ImageAsset
-from vllm.assets.video import VideoAsset
-from vllm.multimodal.image import convert_image_mode
-
-from .conftest import OmniRunner
-from .utils import create_new_process_for_each_test
-
-models = ["Qwen/Qwen2.5-Omni-3B"]
-
-# CI stage config optimized for 24GB GPU (L4/RTX3090)
-CI_STAGE_CONFIG_PATH = str(Path(__file__).parent / "stage_configs" / "qwen2_5_omni_ci.yaml")
-
-
-@pytest.mark.core_model
-@pytest.mark.parametrize("model", models)
-@create_new_process_for_each_test()
-def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
-    """Test processing audio, image, and video together, generating audio output."""
-    with omni_runner(model, seed=42, stage_configs_path=CI_STAGE_CONFIG_PATH) as runner:
-        # Prepare multimodal inputs
-        question = "What is recited in the audio? What is in this image? Describe the video briefly."
-        audio = AudioAsset("mary_had_lamb").audio_and_sample_rate
-        audio = (audio[0][: 16000 * 5], audio[1])  # Trim to first 5 seconds
-        image = convert_image_mode(ImageAsset("cherry_blossom").pil_image.resize((128, 128)), "RGB")
-        video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
-
-        outputs = runner.generate_multimodal(
-            prompts=question,
-            audios=audio,
-            images=image,
-            videos=video,
-        )
-
-        # Verify we got outputs from multiple stages
-        assert len(outputs) > 0
-
-        # Find and verify text output (thinker stage)
-        text_output = None
-        for stage_output in outputs:
-            if stage_output.final_output_type == "text":
-                text_output = stage_output
-                break
-
-        assert text_output is not None
-        assert len(text_output.request_output) > 0
-        text_content = text_output.request_output[0].outputs[0].text
-        assert text_content is not None
-        assert len(text_content.strip()) > 0
-
-        # Find and verify audio output (code2wav stage)
-        audio_output = None
-        for stage_output in outputs:
-            if stage_output.final_output_type == "audio":
-                audio_output = stage_output
-                break
-
-        assert audio_output is not None
-        assert len(audio_output.request_output) > 0
-
-        # Verify audio tensor exists and has content
-        audio_tensor = audio_output.request_output[0].multimodal_output["audio"]
-        assert audio_tensor is not None
-        assert audio_tensor.numel() > 0

tests/multi_stages/utils.py

@@ -1,206 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import contextlib
-import functools
-import os
-import signal
-import subprocess
-import sys
-import tempfile
-from collections.abc import Callable
-from contextlib import ExitStack, suppress
-from pathlib import Path
-from typing import Any, Literal
-
-import cloudpickle
-from typing_extensions import ParamSpec
-from vllm.platforms import current_platform
-
-VLLM_PATH = Path(__file__).parent.parent.parent
-"""Path to root of the vLLM repository."""
-
-
-_P = ParamSpec("_P")
-
-
-def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
-    """Decorator to fork a new process for each test function.
-    See https://github.com/vllm-project/vllm/issues/7053 for more details.
-    """
-
-    @functools.wraps(func)
-    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
-        # Make the process the leader of its own process group
-        # to avoid sending SIGTERM to the parent process
-        os.setpgrp()
-        from _pytest.outcomes import Skipped
-
-        # Create a unique temporary file to store exception info from child
-        # process. Use test function name and process ID to avoid collisions.
-        with (
-            tempfile.NamedTemporaryFile(
-                delete=False,
-                mode="w+b",
-                prefix=f"vllm_test_{func.__name__}_{os.getpid()}_",
-                suffix=".exc",
-            ) as exc_file,
-            ExitStack() as delete_after,
-        ):
-            exc_file_path = exc_file.name
-            delete_after.callback(os.remove, exc_file_path)
-
-            pid = os.fork()
-            print(f"Fork a new process to run a test {pid}")
-            if pid == 0:
-                # Parent process responsible for deleting, don't delete
-                # in child.
-                delete_after.pop_all()
-                try:
-                    func(*args, **kwargs)
-                except Skipped as e:
-                    # convert Skipped to exit code 0
-                    print(str(e))
-                    os._exit(0)
-                except Exception as e:
-                    import traceback
-
-                    tb_string = traceback.format_exc()
-
-                    # Try to serialize the exception object first
-                    exc_to_serialize: dict[str, Any]
-                    try:
-                        # First, try to pickle the actual exception with
-                        # its traceback.
-                        exc_to_serialize = {"pickled_exception": e}
-                        # Test if it can be pickled
-                        cloudpickle.dumps(exc_to_serialize)
-                    except (Exception, KeyboardInterrupt):
-                        # Fall back to string-based approach.
-                        exc_to_serialize = {
-                            "exception_type": type(e).__name__,
-                            "exception_msg": str(e),
-                            "traceback": tb_string,
-                        }
-                    try:
-                        with open(exc_file_path, "wb") as f:
-                            cloudpickle.dump(exc_to_serialize, f)
-                    except Exception:
-                        # Fallback: just print the traceback.
-                        print(tb_string)
-                    os._exit(1)
-                else:
-                    os._exit(0)
-            else:
-                pgid = os.getpgid(pid)
-                _pid, _exitcode = os.waitpid(pid, 0)
-                # ignore SIGTERM signal itself
-                old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
-                # kill all child processes
-                os.killpg(pgid, signal.SIGTERM)
-                # restore the signal handler
-                signal.signal(signal.SIGTERM, old_signal_handler)
-                if _exitcode != 0:
-                    # Try to read the exception from the child process
-                    exc_info = {}
-                    if os.path.exists(exc_file_path):
-                        with (
-                            contextlib.suppress(Exception),
-                            open(exc_file_path, "rb") as f,
-                        ):
-                            exc_info = cloudpickle.load(f)
-
-                    if (original_exception := exc_info.get("pickled_exception")) is not None:
-                        # Re-raise the actual exception object if it was
-                        # successfully pickled.
-                        assert isinstance(original_exception, Exception)
-                        raise original_exception
-
-                    if (original_tb := exc_info.get("traceback")) is not None:
-                        # Use string-based traceback for fallback case
-                        raise AssertionError(
-                            f"Test {func.__name__} failed when called with"
-                            f" args {args} and kwargs {kwargs}"
-                            f" (exit code: {_exitcode}):\n{original_tb}"
-                        ) from None
-
-                    # Fallback to the original generic error
-                    raise AssertionError(
-                        f"function {func.__name__} failed when called with"
-                        f" args {args} and kwargs {kwargs}"
-                        f" (exit code: {_exitcode})"
-                    ) from None
-
-    return wrapper
-
-
-def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
-    """Decorator to spawn a new process for each test function."""
-
-    @functools.wraps(f)
-    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
-        # Check if we're already in a subprocess
-        if os.environ.get("RUNNING_IN_SUBPROCESS") == "1":
-            # If we are, just run the function directly
-            return f(*args, **kwargs)
-
-        import torch.multiprocessing as mp
-
-        with suppress(RuntimeError):
-            mp.set_start_method("spawn")
-
-        # Get the module
-        module_name = f.__module__
-
-        # Create a process with environment variable set
-        env = os.environ.copy()
-        env["RUNNING_IN_SUBPROCESS"] = "1"
-
-        with tempfile.TemporaryDirectory() as tempdir:
-            output_filepath = os.path.join(tempdir, "new_process.tmp")
-
-            # `cloudpickle` allows pickling complex functions directly
-            input_bytes = cloudpickle.dumps((f, output_filepath))
-
-            repo_root = str(VLLM_PATH.resolve())
-
-            env = dict(env or os.environ)
-            env["PYTHONPATH"] = repo_root + os.pathsep + env.get("PYTHONPATH", "")
-
-            cmd = [sys.executable, "-m", f"{module_name}"]
-
-            returned = subprocess.run(cmd, input=input_bytes, capture_output=True, env=env)
-
-            # check if the subprocess is successful
-            try:
-                returned.check_returncode()
-            except Exception as e:
-                # wrap raised exception to provide more information
-                raise RuntimeError(f"Error raised in subprocess:\n{returned.stderr.decode()}") from e
-
-    return wrapper
-
-
-def create_new_process_for_each_test(
-    method: Literal["spawn", "fork"] | None = None,
-) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
-    """Creates a decorator that runs each test function in a new process.
-
-    Args:
-        method: The process creation method. Can be either "spawn" or "fork".
-               If not specified, it defaults to "spawn" on ROCm and XPU
-               platforms and "fork" otherwise.
-
-    Returns:
-        A decorator to run test functions in separate processes.
-    """
-    if method is None:
-        use_spawn = current_platform.is_rocm() or current_platform.is_xpu()
-        method = "spawn" if use_spawn else "fork"
-
-    assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
-
-    if method == "fork":
-        return fork_new_process_for_each_test
-
-    return spawn_new_process_for_each_test

tests/multi_stages_h100/conftest.py

@@ -1,14 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Pytest configuration for multi_stages_h100 tests.
-
-This conftest imports fixtures from the parent multi_stages conftest.
-"""
-
-# Import the omni_runner fixture from the parent directory's conftest
-# This makes the fixture available to tests in this directory
-from ..multi_stages.conftest import omni_runner  # noqa: F401
-
-# Explicitly expose the fixture for pytest discovery
-__all__ = ["omni_runner"]

tests/multi_stages_h100/stage_configs/qwen3_omni_ci.yaml

@@ -1,98 +0,0 @@
-# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
-# Stage 0: Thinker (multimodal understanding + text generation)
-# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
-# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
-
-# The following config has been verified on 2x H100-80G GPUs.
-stage_args:
-  - stage_id: 0
-    runtime:
-      devices: "0,1"
-      max_batch_size: 1
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent  # Output hidden states for talker
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      hf_config_name: thinker_config
-      tensor_parallel_size: 2
-      load_format: dummy
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 100
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-
-  - stage_id: 1
-    runtime:
-       devices: "1"
-       max_batch_size: 1
-    engine_args:
-       model_stage: talker
-       model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.5
-       enforce_eager: true
-       trust_remote_code: true
-       engine_output_type: latent  # Output codec codes for code2wav
-      #  tensor_parallel_size: 2
-       enable_prefix_caching: false
-       distributed_executor_backend: "mp"
-       hf_config_name: talker_config
-       load_format: dummy
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
-    # final_output: true
-    # final_output_type: text
-    default_sampling_params:
-      temperature: 0.9
-      top_k: 50
-      max_tokens: 100
-      seed: 42
-      detokenize: False
-      repetition_penalty: 1.05
-      stop_token_ids: [2150]
-
-  - stage_id: 2
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: code2wav
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      enforce_eager: true
-      trust_remote_code: true
-      enable_prefix_caching: false
-      engine_output_type: audio  # Final output: audio waveform
-      gpu_memory_utilization: 0.1
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 1000000
-      hf_config_name: thinker_config
-      load_format: dummy
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
-    final_output: true
-    final_output_type: audio
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 200
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.1

tests/multi_stages_h100/test_qwen3_omni.py

@@ -1,65 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E tests for Qwen3-Omni model with audio/video/image input and audio output.
-"""
-
-import os
-from pathlib import Path
-
-import pytest
-from vllm.assets.video import VideoAsset
-
-from ..multi_stages.conftest import OmniRunner
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
-
-# CI stage config for 2xH100-80G GPUs
-CI_STAGE_CONFIG_PATH = str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")
-
-
-@pytest.mark.parametrize("model", models)
-def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
-    """Test processing video, generating audio output."""
-    with omni_runner(model, seed=42, stage_configs_path=CI_STAGE_CONFIG_PATH) as runner:
-        # Prepare inputs
-        question = "Describe the video briefly."
-        video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
-
-        outputs = runner.generate_multimodal(
-            prompts=question,
-            videos=video,
-        )
-
-        # Verify we got outputs from multiple stages
-        assert len(outputs) > 0
-
-        # Find and verify text output (thinker stage)
-        text_output = None
-        for stage_output in outputs:
-            if stage_output.final_output_type == "text":
-                text_output = stage_output
-                break
-
-        assert text_output is not None
-        assert len(text_output.request_output) > 0
-        text_content = text_output.request_output[0].outputs[0].text
-        assert text_content is not None
-        assert len(text_content.strip()) > 0
-
-        # Find and verify audio output (code2wav stage)
-        audio_output = None
-        for stage_output in outputs:
-            if stage_output.final_output_type == "audio":
-                audio_output = stage_output
-                break
-
-        assert audio_output is not None
-        assert len(audio_output.request_output) > 0
-
-        # Verify audio tensor exists and has content
-        audio_tensor = audio_output.request_output[0].multimodal_output["audio"]
-        assert audio_tensor is not None
-        assert audio_tensor.numel() > 0

tests/single_stage/test_diffusion_model.py

@@ -1,39 +0,0 @@
-import os
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-# ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[2]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from vllm_omni import Omni
-
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
-
-models = ["Tongyi-MAI/Z-Image-Turbo", "riverclouds/qwen_image_random"]
-
-
-@pytest.mark.parametrize("model_name", models)
-def test_diffusion_model(model_name: str):
-    m = Omni(model=model_name)
-    # high resolution may cause OOM on L4
-    height = 256
-    width = 256
-    images = m.generate(
-        "a photo of a cat sitting on a laptop keyboard",
-        height=height,
-        width=width,
-        num_inference_steps=2,
-        guidance_scale=0.0,
-        generator=torch.Generator("cuda").manual_seed(42),
-        num_outputs_per_prompt=2,
-    )
-    assert len(images) == 2
-    # check image size
-    assert images[0].width == width
-    assert images[0].height == height
-    images[0].save("image_output.png")

tests/test_omni_llm.py

@@ -342,14 +342,14 @@ def _setup_log_mocks(monkeypatch):
 def mock_get_config(monkeypatch):
     """Auto-mock get_config and related model loading functions to avoid model path validation."""
     # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports
-    # This prevents ImportError when async_omni is imported (which happens via omni_stage)
+    # This prevents ImportError when async_omni_llm is imported (which happens via omni_stage)
     import sys
 
     fake_tokenizer = MagicMock()
     fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3])
     fake_tokenizer.decode = MagicMock(return_value="test")
 
-    # Mock init_tokenizer_from_configs (used in async_omni)
+    # Mock init_tokenizer_from_configs (used in async_omni_llm)
     def _mock_init_tokenizer_from_configs(model_config=None, **kwargs):
         return fake_tokenizer
 
@@ -367,9 +367,8 @@ def mock_get_config(monkeypatch):
         tokenizer_module = sys.modules[tokenizer_module_path]
         setattr(tokenizer_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
 
-    # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni
-
-    # This is because async_omni imports processor.py, which imports this function at module level
+    # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni_llm
+    # This is because async_omni_llm imports processor.py, which imports this function at module level
     # Mock length_from_prompt_token_ids_or_embeds (used in processor.py)
     def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None):
         # Return a reasonable default length
@@ -403,19 +402,19 @@ def mock_get_config(monkeypatch):
             processor_module, "length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds
         )
 
-    # Strategy 3: Now mock async_omni AFTER length_from_prompt_token_ids_or_embeds is mocked
-    # This prevents ImportError when async_omni imports processor.py
+    # Strategy 3: Now mock async_omni_llm AFTER length_from_prompt_token_ids_or_embeds is mocked
+    # This prevents ImportError when async_omni_llm imports processor.py
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.async_omni.init_tokenizer_from_configs",
+        "vllm_omni.entrypoints.async_omni_llm.init_tokenizer_from_configs",
         _mock_init_tokenizer_from_configs,
         raising=False,
     )
 
-    # Strategy 4: If async_omni is already imported, patch it directly
-    async_omni_path = "vllm_omni.entrypoints.async_omni"
-    if async_omni_path in sys.modules:
-        async_omni_module = sys.modules[async_omni_path]
-        setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
+    # Strategy 4: If async_omni_llm is already imported, patch it directly
+    async_omni_llm_path = "vllm_omni.entrypoints.async_omni_llm"
+    if async_omni_llm_path in sys.modules:
+        async_omni_llm_module = sys.modules[async_omni_llm_path]
+        setattr(async_omni_llm_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
 
     # Now mock get_config and other functions
     fake_hf_config = MagicMock()

tests/utils.py

@@ -1,107 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import os
-import time
-from contextlib import contextmanager
-
-from vllm.platforms import current_platform
-
-if current_platform.is_rocm():
-    from amdsmi import (
-        amdsmi_get_gpu_vram_usage,
-        amdsmi_get_processor_handles,
-        amdsmi_init,
-        amdsmi_shut_down,
-    )
-
-    @contextmanager
-    def _nvml():
-        try:
-            amdsmi_init()
-            yield
-        finally:
-            amdsmi_shut_down()
-elif current_platform.is_cuda():
-    from vllm.third_party.pynvml import (
-        nvmlDeviceGetHandleByIndex,
-        nvmlDeviceGetMemoryInfo,
-        nvmlInit,
-        nvmlShutdown,
-    )
-
-    @contextmanager
-    def _nvml():
-        try:
-            nvmlInit()
-            yield
-        finally:
-            nvmlShutdown()
-else:
-
-    @contextmanager
-    def _nvml():
-        yield
-
-
-def get_physical_device_indices(devices):
-    visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
-    if visible_devices is None:
-        return devices
-
-    visible_indices = [int(x) for x in visible_devices.split(",")]
-    index_mapping = {i: physical for i, physical in enumerate(visible_indices)}
-    return [index_mapping[i] for i in devices if i in index_mapping]
-
-
-@_nvml()
-def wait_for_gpu_memory_to_clear(
-    *,
-    devices: list[int],
-    threshold_bytes: int | None = None,
-    threshold_ratio: float | None = None,
-    timeout_s: float = 120,
-) -> None:
-    assert threshold_bytes is not None or threshold_ratio is not None
-    # Use nvml instead of pytorch to reduce measurement error from torch cuda
-    # context.
-    devices = get_physical_device_indices(devices)
-    start_time = time.time()
-    while True:
-        output: dict[int, str] = {}
-        output_raw: dict[int, tuple[float, float]] = {}
-        for device in devices:
-            if current_platform.is_rocm():
-                dev_handle = amdsmi_get_processor_handles()[device]
-                mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
-                gb_used = mem_info["vram_used"] / 2**10
-                gb_total = mem_info["vram_total"] / 2**10
-            else:
-                dev_handle = nvmlDeviceGetHandleByIndex(device)
-                mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
-                gb_used = mem_info.used / 2**30
-                gb_total = mem_info.total / 2**30
-            output_raw[device] = (gb_used, gb_total)
-            output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
-
-        print("gpu memory used/total (GiB): ", end="")
-        for k, v in output.items():
-            print(f"{k}={v}; ", end="")
-        print("")
-
-        if threshold_bytes is not None:
-            is_free = lambda used, total: used <= threshold_bytes / 2**30  # noqa E731
-            threshold = f"{threshold_bytes / 2**30} GiB"
-        else:
-            is_free = lambda used, total: used / total <= threshold_ratio  # noqa E731
-            threshold = f"{threshold_ratio:.2f}"
-
-        dur_s = time.time() - start_time
-        if all(is_free(used, total) for used, total in output_raw.values()):
-            print(f"Done waiting for free GPU memory on devices {devices=} ({threshold=}) {dur_s=:.02f}")
-            break
-
-        if dur_s >= timeout_s:
-            raise ValueError(f"Memory of devices {devices=} not free after {dur_s=:.02f} ({threshold=})")
-
-        time.sleep(5)

vllm_omni/__init__.py

@@ -21,10 +21,10 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
     patch = None  # type: ignore
 
 from .config import OmniModelConfig
-from .entrypoints.async_omni import AsyncOmni
+from .entrypoints.async_omni_llm import AsyncOmniLLM
 
 # Main entry points
-from .entrypoints.omni import Omni
+# from .entrypoints.omni import Omni
 
 from .version import __version__, __version_tuple__  # isort:skip
 
@@ -33,8 +33,8 @@ __all__ = [
     "__version__",
     "__version_tuple__",
     # Main components
-    "Omni",
-    "AsyncOmni",
+    # "Omni",
+    "AsyncOmniLLM",
     # Configuration
     "OmniModelConfig",
     # All other components are available through their respective modules

vllm_omni/core/sched/__init__.py

@@ -2,12 +2,12 @@
 Scheduling components for vLLM-Omni.
 """
 
-from .omni_ar_scheduler import OmniARScheduler
-from .omni_generation_scheduler import OmniGenerationScheduler
+from .diffusion_scheduler import DiffusionScheduler
 from .output import OmniNewRequestData
+from .scheduler import OmniScheduler
 
 __all__ = [
-    "OmniARScheduler",
-    "OmniGenerationScheduler",
+    "OmniScheduler",
+    "DiffusionScheduler",
     "OmniNewRequestData",
 ]

vllm_omni/core/sched/diffusion_scheduler.py

@@ -0,0 +1,315 @@
+import time
+from collections import defaultdict
+from typing import Optional
+
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.sched.request_queue import create_request_queue
+from vllm.v1.core.sched.scheduler import EngineCoreOutputs, Request, RequestStatus, SchedulerOutput, SpecDecodingStats
+from vllm.v1.core.sched.utils import remove_all
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
+
+from vllm_omni.core.sched.output import OmniNewRequestData
+from vllm_omni.core.sched.scheduler import OmniScheduler
+from vllm_omni.outputs import OmniModelRunnerOutput
+
+
+class DiffusionScheduler(OmniScheduler):
+    """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:
+        """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
+        # (diffusion results returned via pooler_output)
+        token_budget = self.max_num_scheduled_tokens
+        capacity = self.max_num_running_reqs - len(self.running)
+        scheduled_timestamp = time.monotonic()
+
+        scheduled_new_reqs: list[Request] = []
+        scheduled_resumed_reqs: list[Request] = []
+        scheduled_running_reqs: list[Request] = []
+
+        req_to_new_blocks: dict[str, KVCacheBlocks] = {}
+        num_scheduled_tokens: dict[str, int] = {}
+        scheduled_spec_decode_tokens: dict[str, list[int]] = {}
+        scheduled_encoder_inputs: dict[str, list[int]] = {}
+
+        # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
+        skipped_waiting_requests = create_request_queue(self.policy)
+
+        # Fast path selection and scheduling (treat all as diffusion requests,
+        # independent of pooling_params)
+        while self.waiting and token_budget > 0 and capacity > 0:
+            request = self.waiting.peek_request()
+            # Uniformly treat as diffusion. A feature flag can be added later
+            # via config or request tag.
+
+            # Allocate all input tokens for the request in one shot
+            # (allocate 1 placeholder if zero)
+            required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
+            if required_tokens > token_budget:
+                # Insufficient budget to process all inputs at once;
+                # stop fast path attempt
+                break
+            num_new_tokens = required_tokens
+            new_blocks = self.kv_cache_manager.allocate_slots(
+                request,
+                num_new_tokens,
+                num_lookahead_tokens=self.num_lookahead_tokens,
+            )
+            if new_blocks is None:
+                # Allocation failed (e.g., VRAM pressure); stop fast path and
+                # fall back to default scheduling
+                # Put the current request back to the head of the waiting queue
+                # Note: the original queue order is preserved
+                break
+
+            # Officially schedule this request
+            request = self.waiting.pop_request()
+            self.running.append(request)
+            request.status = RequestStatus.RUNNING
+            if self.log_stats:
+                request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
+
+            req_to_new_blocks[request.request_id] = new_blocks
+            num_scheduled_tokens[request.request_id] = num_new_tokens
+            token_budget -= num_new_tokens
+            capacity -= 1
+            scheduled_new_reqs.append(request)
+
+        # Return skipped waiting requests
+        if skipped_waiting_requests:
+            self.waiting.prepend_requests(skipped_waiting_requests)
+
+        # If fast path scheduled none, fall back to the original scheduling
+        if not num_scheduled_tokens:
+            return super().schedule()
+
+        # Compute common prefix blocks (aligned with v1)
+        num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
+        if self.running:
+            any_request = self.running[0]
+            num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(
+                any_request, len(self.running)
+            )
+
+        # Assemble SchedulerOutput
+        new_reqs_data = [
+            OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids())
+            for req in scheduled_new_reqs
+        ]
+        cached_reqs_data = self._make_cached_request_data(
+            scheduled_running_reqs,
+            scheduled_resumed_reqs,
+            num_scheduled_tokens,
+            scheduled_spec_decode_tokens,
+            req_to_new_blocks,
+        )
+        scheduled_requests = scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
+        structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(
+            scheduled_requests, scheduled_spec_decode_tokens
+        )
+
+        total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=new_reqs_data,
+            scheduled_cached_reqs=cached_reqs_data,
+            num_scheduled_tokens=num_scheduled_tokens,
+            total_num_scheduled_tokens=total_num_scheduled_tokens,
+            scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
+            scheduled_encoder_inputs=scheduled_encoder_inputs,
+            num_common_prefix_blocks=num_common_prefix_blocks,
+            finished_req_ids=self.finished_req_ids,
+            free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
+            structured_output_request_ids=structured_output_request_ids,
+            grammar_bitmask=grammar_bitmask,
+        )
+
+        # KVTransfer: package metadata
+        if self.connector is not None:
+            meta = self.connector.build_connector_meta(scheduler_output)
+            scheduler_output.kv_connector_metadata = meta
+
+        # Update internal state (advance num_computed_tokens, free encoder inputs,
+        # etc.)
+        self._update_after_schedule(scheduler_output)
+        return scheduler_output
+
+    """
+    Scheduler for the diffusion model.
+    This scheduler is modified to stop the request immediately for the diffusion model.
+    This is because the diffusion model can generate the final image/audio in one step.
+    Note: This is just a minimal modification to the original scheduler,
+    and there should be some further efforts to optimize the scheduler.
+    The original scheduler is still used for the AR model.
+    """
+
+    def update_from_output(
+        self,
+        scheduler_output: SchedulerOutput,
+        model_runner_output: OmniModelRunnerOutput,
+    ) -> dict[int, EngineCoreOutputs]:
+        """Update the scheduler state based on the model runner output.
+
+        This method is modified to stop the request immediately for the diffusion model.
+        """
+        sampled_token_ids = model_runner_output.sampled_token_ids
+        logprobs = model_runner_output.logprobs
+        prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
+        num_scheduled_tokens = scheduler_output.num_scheduled_tokens
+        pooler_outputs = model_runner_output.pooler_output
+        num_nans_in_logits = model_runner_output.num_nans_in_logits
+        kv_connector_output = model_runner_output.kv_connector_output
+
+        outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
+        spec_decoding_stats: Optional[SpecDecodingStats] = None
+        kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
+
+        # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
+        # the below loop can be a performance bottleneck. We should do our best
+        # to avoid expensive operations inside the loop.
+        stopped_running_reqs: set[Request] = set()
+        stopped_preempted_reqs: set[Request] = set()
+        for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
+            assert num_tokens_scheduled > 0
+            request = self.requests.get(req_id)
+            if request is None:
+                # The request is already finished. This can happen if the
+                # request is aborted while the model is executing it (e.g.,
+                # in pipeline parallelism).
+                continue
+
+            req_index = model_runner_output.req_id_to_index[req_id]
+            generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
+
+            scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
+            if scheduled_spec_token_ids:
+                num_draft_tokens = len(scheduled_spec_token_ids)
+                num_accepted = len(generated_token_ids) - 1
+                num_rejected = num_draft_tokens - num_accepted
+                # num_computed_tokens represents the number of tokens
+                # processed in the current step, considering scheduled
+                # tokens and rejections. If some tokens are rejected,
+                # num_computed_tokens is decreased by the number of rejected
+                # tokens.
+                request.num_computed_tokens -= num_rejected
+                spec_decoding_stats = self.make_spec_decoding_stats(
+                    spec_decoding_stats,
+                    num_draft_tokens=num_draft_tokens,
+                    num_accepted_tokens=num_accepted,
+                )
+
+            new_logprobs = None
+            new_token_ids = generated_token_ids
+            kv_transfer_params = None
+            status_before_stop = request.status
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]
+
+            # Diffusion request: completes in one step; mark finished and free resources
+            request.status = RequestStatus.FINISHED_STOPPED
+            # Optional: set a stop_reason for front-end clarity
+            # (does not affect protocol)
+            request.stop_reason = request.stop_reason or "diffusion_done"
+            kv_transfer_params = self._free_request(request)
+            if status_before_stop == RequestStatus.RUNNING:
+                stopped_running_reqs.add(request)
+            else:
+                stopped_preempted_reqs.add(request)
+
+            # Extract sample logprobs if needed.
+            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
+                # NOTE: once we support N tokens per step (spec decode),
+                # the outer lists can be of length > 1.
+                new_logprobs = logprobs.slice(req_index, req_index + 1)
+
+            if new_token_ids and self.structured_output_manager.should_advance(request):
+                # NOTE: structured_output_request should not be None if
+                # use_structured_output, we have check above, so safe to ignore
+                # type warning
+                request.structured_output_request.grammar.accept_tokens(req_id, new_token_ids)  # type: ignore[union-attr]  # noqa: E501
+
+            # spec_token_ids comes from the model runner output
+            if num_nans_in_logits is not None and req_id in num_nans_in_logits:
+                request.num_nans_in_logits = num_nans_in_logits[req_id]
+
+            # Get prompt logprobs for this request.
+            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
+            # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput
+            pooling_output_dict = None
+            if pooler_output is not None:
+                # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type
+                # which expects Optional[dict[str, torch.Tensor]]
+                pooling_output_dict = {"model_outputs": pooler_output}
+            if new_token_ids or pooler_output is not None or kv_transfer_params:
+                # Add EngineCoreOutput for this Request.
+                outputs[request.client_index].append(
+                    EngineCoreOutput(
+                        request_id=req_id,
+                        new_token_ids=new_token_ids,
+                        finish_reason=request.get_finished_reason(),
+                        new_logprobs=new_logprobs,
+                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                        pooling_output=pooling_output_dict,
+                        stop_reason=request.stop_reason,
+                        events=request.take_events(),
+                        kv_transfer_params=kv_transfer_params,
+                        trace_headers=request.trace_headers,
+                        num_cached_tokens=request.num_cached_tokens,
+                    )
+                )
+            else:
+                # Invariant: EngineCore returns no partial prefill outputs.
+                assert not prompt_logprobs_tensors
+
+        # Remove the stopped requests from the running and waiting queues.
+        if stopped_running_reqs:
+            if stopped_running_reqs:
+                self.running = remove_all(self.running, stopped_running_reqs)
+        if stopped_preempted_reqs:
+            # This is a rare case and unlikely to impact performance.
+            self.waiting.remove_requests(stopped_preempted_reqs)
+
+        # KV Connector: update state for finished KV Transfers.
+        if model_runner_output.kv_connector_output:
+            self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output)
+
+        # Create EngineCoreOutputs for all clients that have requests with
+        # outputs in this step.
+        engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
+
+        finished_req_ids = self.finished_req_ids_dict
+        if finished_req_ids:
+            # Include ids of requests that finished since last outputs
+            # were sent.
+            for client_index, finished_set in finished_req_ids.items():
+                # Set finished request set in EngineCoreOutputs for this client.
+                if (eco := engine_core_outputs.get(client_index)) is not None:
+                    eco.finished_requests = finished_set
+                else:
+                    engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
+            finished_req_ids.clear()
+
+        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
+            # Return stats to only one of the front-ends.
+            if (eco := next(iter(engine_core_outputs.values()), None)) is None:
+                # We must return the stats even if there are no request
+                # outputs this step.
+                engine_core_outputs[0] = eco = EngineCoreOutputs()
+            eco.scheduler_stats = stats
+
+        return engine_core_outputs

vllm_omni/core/sched/generation_scheduler.py

@@ -5,20 +5,21 @@ from typing import Optional
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.request_queue import create_request_queue
 from vllm.v1.core.sched.scheduler import (
+    EngineCoreOutputs,
     Request,
     RequestStatus,
     SchedulerOutput,
     SpecDecodingStats,
 )
-from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
 
 from vllm_omni.core.sched.output import OmniNewRequestData
+from vllm_omni.core.sched.scheduler import OmniScheduler
 from vllm_omni.outputs import OmniModelRunnerOutput
 
 
-class OmniGenerationScheduler(VLLMScheduler):
+class GenerationScheduler(OmniScheduler):
     def schedule(self) -> SchedulerOutput:
         """Diffusion fast path:
         - Feed all input tokens of the request at once

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -1,233 +0,0 @@
-from __future__ import annotations
-
-from collections import defaultdict
-
-from vllm.logger import init_logger
-from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
-from vllm.v1.core.sched.utils import check_stop, remove_all
-from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
-from vllm.v1.outputs import ModelRunnerOutput
-from vllm.v1.request import Request, RequestStatus
-from vllm.v1.spec_decode.metrics import SpecDecodingStats
-
-
-class OmniARScheduler(VLLMScheduler):
-    """
-    OmniScheduler: Scheduler for vLLM-Omni multimodal processing.
-
-    This scheduler extends vLLM's scheduler to support multimodal and
-    non-autoregressive processing with additional fields and methods
-    specific to vLLM-Omni.
-    """
-
-    # Ensure scheduled_new_reqs carry omni-specific payloads
-    # (e.g., additional_information)
-    def schedule(self) -> SchedulerOutput:  # type: ignore[override]
-        scheduler_output = super().schedule()
-        try:
-            # Late import to avoid circulars in some launch modes
-            from .output import OmniNewRequestData
-
-            # Rewrap base NewRequestData entries with OmniNewRequestData,
-            # enriching with request-level payloads
-            new_list = []
-            for nr in scheduler_output.scheduled_new_reqs:
-                req_id = getattr(nr, "req_id", None)
-                request = self.requests.get(req_id) if req_id else None
-                # Build omni entry preserving all base fields
-                omni_nr = OmniNewRequestData(
-                    req_id=nr.req_id,
-                    prompt_token_ids=nr.prompt_token_ids,
-                    mm_features=nr.mm_features,
-                    sampling_params=nr.sampling_params,
-                    pooling_params=nr.pooling_params,
-                    block_ids=nr.block_ids,
-                    num_computed_tokens=nr.num_computed_tokens,
-                    lora_request=nr.lora_request,
-                    # Enrich with omni payloads from the live request object
-                    prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
-                    additional_information=(getattr(request, "additional_information", None) if request else None),
-                )
-                new_list.append(omni_nr)
-
-            scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
-        except Exception:
-            # If anything goes wrong, leave the original output unchanged
-            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
-
-        return scheduler_output
-
-    def update_from_output(
-        self,
-        scheduler_output: SchedulerOutput,
-        model_runner_output: ModelRunnerOutput,
-    ) -> dict[int, EngineCoreOutputs]:
-        """Override update_from_output to fix check_stop bug when output_token_ids is empty.
-
-        The original vLLM implementation calls check_stop when pooler_outputs exist,
-        but doesn't check if output_token_ids is empty. This causes IndexError during
-        prefill phase when pooler_outputs exist but output_token_ids is still empty.
-
-        This override adds a safety check before calling check_stop.
-        """
-        sampled_token_ids = model_runner_output.sampled_token_ids
-        logprobs = model_runner_output.logprobs
-        prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
-        num_scheduled_tokens = scheduler_output.num_scheduled_tokens
-        pooler_outputs = model_runner_output.pooler_output
-        num_nans_in_logits = model_runner_output.num_nans_in_logits
-        kv_connector_output = model_runner_output.kv_connector_output
-
-        outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
-        spec_decoding_stats: SpecDecodingStats | None = None
-        kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
-
-        # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
-        # the below loop can be a performance bottleneck. We should do our best
-        # to avoid expensive operations inside the loop.
-        stopped_running_reqs: set[Request] = set()
-        stopped_preempted_reqs: set[Request] = set()
-        for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
-            assert num_tokens_scheduled > 0
-            request = self.requests.get(req_id)
-            if request is None:
-                # The request is already finished. This can happen if the
-                # request is aborted while the model is executing it (e.g.,
-                # in pipeline parallelism).
-                continue
-
-            req_index = model_runner_output.req_id_to_index[req_id]
-            generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
-
-            scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
-            if scheduled_spec_token_ids:
-                num_draft_tokens = len(scheduled_spec_token_ids)
-                num_accepted = len(generated_token_ids) - 1
-                num_rejected = num_draft_tokens - num_accepted
-                # num_computed_tokens represents the number of tokens
-                # processed in the current step, considering scheduled
-                # tokens and rejections. If some tokens are rejected,
-                # num_computed_tokens is decreased by the number of rejected
-                # tokens.
-                request.num_computed_tokens -= num_rejected
-                spec_decoding_stats = self.make_spec_decoding_stats(
-                    spec_decoding_stats, num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted
-                )
-
-            stopped = False
-            new_logprobs = None
-            new_token_ids = generated_token_ids
-            kv_transfer_params = None
-            status_before_stop = request.status
-
-            # Check for stop and update request status.
-            if new_token_ids:
-                new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
-
-            """
-            # Stop checking for pooler models.
-            pooler_output = None
-            if pooler_outputs:
-                pooler_output = pooler_outputs[req_index]
-                stopped = check_stop(request, self.max_model_len,
-                                     pooler_output)
-            """
-
-            # Stop checking for pooler models.
-            # FIX: Check if output_token_ids is empty before calling check_stop
-            # to prevent IndexError during prefill phase when pooler_outputs exist
-            # but output_token_ids is still empty.
-            pooler_output = None
-            if pooler_outputs:
-                pooler_output = pooler_outputs[req_index]
-                # Only call check_stop if output_token_ids is not empty
-                # This prevents IndexError when accessing output_token_ids[-1]
-                if request.output_token_ids:
-                    stopped = check_stop(request, self.max_model_len, pooler_output)
-                # If output_token_ids is empty, we don't call check_stop
-                # and keep stopped=False (don't stop during prefill phase)
-
-            if stopped:
-                kv_transfer_params = self._free_request(request)
-                if status_before_stop == RequestStatus.RUNNING:
-                    stopped_running_reqs.add(request)
-                else:
-                    stopped_preempted_reqs.add(request)
-
-            # Extract sample logprobs if needed.
-            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
-                # NOTE: once we support N tokens per step (spec decode),
-                # the outer lists can be of length > 1.
-                new_logprobs = logprobs.slice(req_index, req_index + 1)
-
-            if new_token_ids and self.structured_output_manager.should_advance(request):
-                # NOTE: structured_output_request
-                # should not be None if use_structured_output, we have
-                # checked above, so safe to ignore type warning
-                request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]
-                    req_id, new_token_ids
-                )
-
-            if num_nans_in_logits is not None and req_id in num_nans_in_logits:
-                request.num_nans_in_logits = num_nans_in_logits[req_id]
-
-            # Get prompt logprobs for this request.
-            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or pooler_output is not None or kv_transfer_params:
-                # Add EngineCoreOutput for this Request.
-                outputs[request.client_index].append(
-                    EngineCoreOutput(
-                        request_id=req_id,
-                        new_token_ids=new_token_ids,
-                        finish_reason=request.get_finished_reason(),
-                        new_logprobs=new_logprobs,
-                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                        pooling_output=pooler_output,
-                        stop_reason=request.stop_reason,
-                        events=request.take_events(),
-                        kv_transfer_params=kv_transfer_params,
-                        trace_headers=request.trace_headers,
-                        num_cached_tokens=request.num_cached_tokens,
-                    )
-                )
-            else:
-                # Invariant: EngineCore returns no partial prefill outputs.
-                assert not prompt_logprobs_tensors
-
-        # Remove the stopped requests from the running and waiting queues.
-        if stopped_running_reqs:
-            self.running = remove_all(self.running, stopped_running_reqs)
-        if stopped_preempted_reqs:
-            # This is a rare case and unlikely to impact performance.
-            self.waiting.remove_requests(stopped_preempted_reqs)
-
-        # KV Connector: update state for finished KV Transfers.
-        if model_runner_output.kv_connector_output:
-            self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output)
-
-        # Create EngineCoreOutputs for all clients that have requests with
-        # outputs in this step.
-        engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
-
-        finished_req_ids = self.finished_req_ids_dict
-        if finished_req_ids:
-            # Include ids of requests that finished since last outputs
-            # were sent.
-            for client_index, finished_set in finished_req_ids.items():
-                # Set finished request set in EngineCoreOutputs for this client.
-                if (eco := engine_core_outputs.get(client_index)) is not None:
-                    eco.finished_requests = finished_set
-                else:
-                    engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
-            finished_req_ids.clear()
-
-        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
-            # Return stats to only one of the front-ends.
-            if (eco := next(iter(engine_core_outputs.values()), None)) is None:
-                # We must return the stats even if there are no request
-                # outputs this step.
-                engine_core_outputs[0] = eco = EngineCoreOutputs()
-            eco.scheduler_stats = stats
-
-        return engine_core_outputs

vllm_omni/core/sched/scheduler.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
+
+
+class OmniScheduler(VLLMScheduler):
+    """
+    OmniScheduler: Scheduler for vLLM-Omni multimodal processing.
+
+    This scheduler extends vLLM's scheduler to support multimodal and
+    non-autoregressive processing with additional fields and methods
+    specific to vLLM-Omni.
+    """
+
+    # Ensure scheduled_new_reqs carry omni-specific payloads
+    # (e.g., additional_information)
+    def schedule(self) -> SchedulerOutput:  # type: ignore[override]
+        scheduler_output = super().schedule()
+        try:
+            # Late import to avoid circulars in some launch modes
+            from .output import OmniNewRequestData
+
+            # Rewrap base NewRequestData entries with OmniNewRequestData,
+            # enriching with request-level payloads
+            new_list = []
+            for nr in scheduler_output.scheduled_new_reqs:
+                req_id = getattr(nr, "req_id", None)
+                request = self.requests.get(req_id) if req_id else None
+                # Build omni entry preserving all base fields
+                omni_nr = OmniNewRequestData(
+                    req_id=nr.req_id,
+                    prompt_token_ids=nr.prompt_token_ids,
+                    mm_features=nr.mm_features,
+                    sampling_params=nr.sampling_params,
+                    pooling_params=nr.pooling_params,
+                    block_ids=nr.block_ids,
+                    num_computed_tokens=nr.num_computed_tokens,
+                    lora_request=nr.lora_request,
+                    # Enrich with omni payloads from the live request object
+                    prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
+                    additional_information=(getattr(request, "additional_information", None) if request else None),
+                )
+                new_list.append(omni_nr)
+
+            scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
+        except Exception:
+            # If anything goes wrong, leave the original output unchanged
+            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
+
+        return scheduler_output

vllm_omni/diffusion/attention/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/diffusion/attention/backends/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/diffusion/attention/backends/abstract.py

@@ -1,78 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from typing import Generic, TypeVar
-
-import torch
-
-
-class AttentionBackend(ABC):
-    """Abstract class for diffusion attention backends."""
-
-    accept_output_buffer: bool = False
-
-    @staticmethod
-    @abstractmethod
-    def get_name() -> str:
-        raise NotImplementedError
-
-    @staticmethod
-    @abstractmethod
-    def get_impl_cls() -> type["AttentionImpl"]:
-        raise NotImplementedError
-
-    @staticmethod
-    @abstractmethod
-    def get_metadata_cls() -> type["AttentionMetadata"]:
-        raise NotImplementedError
-
-    @staticmethod
-    @abstractmethod
-    def get_builder_cls():  # -> Type["AttentionMetadataBuilder"]:
-        raise NotImplementedError
-
-    @staticmethod
-    @abstractmethod
-    def get_supported_head_sizes() -> list[int]:
-        """Get the list of supported head sizes for this backend."""
-        raise NotImplementedError
-
-    @classmethod
-    def supports_head_size(cls, head_size: int) -> bool:
-        supported_head_sizes = cls.get_supported_head_sizes()
-        return (not supported_head_sizes) or head_size in supported_head_sizes
-
-
-@dataclass
-class AttentionMetadata:
-    attn_mask: torch.Tensor | None = None
-
-
-T = TypeVar("T", bound=AttentionMetadata)
-
-
-class AttentionImpl(ABC, Generic[T]):
-    @abstractmethod
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        softmax_scale: float,
-        causal: bool = False,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        raise NotImplementedError
-
-    @abstractmethod
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: T,
-    ) -> torch.Tensor:
-        raise NotImplementedError

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -1,74 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.attention.backends.abstract import (
-    AttentionBackend,
-    AttentionImpl,
-    AttentionMetadata,
-)
-
-logger = init_logger(__name__)
-
-try:
-    # only tested with flash_attn v3
-    # from flash_attn_interface import flash_attn_func as flash_attn_3_func  # not available in flash-attn 2.8.1
-    from flash_attn import flash_attn_func  # can be FA2 or FA3
-except ImportError:
-    logger.warning(
-        "FlashAttentionBackend is not available. You may install flash-attn "
-        "by running `uv pip install flash-attn==2.8.1 --no-build-isolation`"
-        " or install pre-built flash-attn from https://github.com/Dao-AILab/flash-attention/releases"
-    )
-    flash_attn_3_func = None
-
-
-class FlashAttentionBackend(AttentionBackend):
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        return [64, 96, 128, 192, 256]
-
-    @staticmethod
-    def get_name() -> str:
-        return "FLASH_ATTN"
-
-    @staticmethod
-    def get_impl_cls() -> type["FlashAttentionImpl"]:
-        return FlashAttentionImpl
-
-
-class FlashAttentionImpl(AttentionImpl):
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        softmax_scale: float,
-        causal: bool = False,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        self.num_heads = num_heads
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> torch.Tensor:
-        # TODO: flash_attn_func does not support attn_mask.
-        out: torch.Tensor = flash_attn_func(
-            query,
-            key,
-            value,
-            causal=self.causal,
-            softmax_scale=self.softmax_scale,
-        )
-        return out

vllm_omni/diffusion/attention/backends/sdpa.py

@@ -1,66 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.attention.backends.abstract import (
-    AttentionBackend,
-    AttentionImpl,
-    AttentionMetadata,
-)
-
-logger = init_logger(__name__)
-
-
-class SDPABackend(AttentionBackend):
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        return [x for x in range(1024)]  # todo
-
-    @staticmethod
-    def get_name() -> str:
-        return "SDPA"
-
-    @staticmethod
-    def get_impl_cls() -> type["SDPAImpl"]:
-        return SDPAImpl
-
-
-class SDPAImpl(AttentionImpl):
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        softmax_scale: float,
-        causal: bool = False,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> torch.Tensor:
-        query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
-        attention_mask = attn_metadata.attn_mask if attn_metadata else None
-
-        output = torch.nn.functional.scaled_dot_product_attention(
-            query,
-            key,
-            value,
-            attn_mask=attention_mask,
-            dropout_p=0.0,
-            is_causal=self.causal,
-            scale=self.softmax_scale,
-        )
-        out = output.permute(0, 2, 1, 3)
-        return out

vllm_omni/diffusion/attention/layer.py

@@ -1,43 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-import torch.nn as nn
-
-from vllm_omni.diffusion.attention.backends.abstract import (
-    AttentionMetadata,
-)
-from vllm_omni.diffusion.attention.selector import get_attn_backend
-
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        causal: bool,
-        softmax_scale: float,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.attn_backend = get_attn_backend(-1)
-        self.attn_impl_cls = self.attn_backend.get_impl_cls()
-        self.attention = self.attn_impl_cls(
-            num_heads=num_heads,
-            head_size=head_size,
-            softmax_scale=softmax_scale,
-            causal=causal,
-            num_kv_heads=num_kv_heads,
-        )
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> torch.Tensor:
-        # shape: (batch_size, seq_len, num_heads, head_size)
-        attn_output = self.attention.forward(query, key, value, attn_metadata)
-        return attn_output

vllm_omni/diffusion/attention/selector.py

@@ -1,53 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import os
-from functools import cache
-
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.attention.backends.abstract import (
-    AttentionBackend,
-)
-from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
-from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
-
-logger = init_logger(__name__)
-
-# environment variable value -> backend class
-SUPPORTED_BACKENDS = {
-    "FLASH_ATTN": FlashAttentionBackend,
-    "TORCH_SDPA": SDPABackend,
-}
-
-
-@cache
-def get_attn_backend(head_size: int) -> type[AttentionBackend]:
-    """
-    Get attention backend for diffusion models.
-
-    The backend is selected based on the following priority:
-    1. DIFFUSION_ATTENTION_BACKEND environment variable (if set, e.g. export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN)
-    2. Default backend (SDPA)
-
-    Args:
-        head_size: Head size (currently not used for selection, but kept for API compatibility)
-
-    Returns:
-        The selected attention backend class
-    """
-    # Check environment variable
-    backend_name: str | None = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
-
-    if backend_name is not None:
-        backend_name_upper = backend_name.upper()
-        if backend_name_upper not in SUPPORTED_BACKENDS:
-            valid_backends = list(SUPPORTED_BACKENDS.keys())
-            raise ValueError(
-                f"Invalid attention backend for diffusion: '{backend_name}'. Valid backends are: {valid_backends}"
-            )
-        logger.info(f"Using attention backend '{backend_name_upper}' for diffusion")
-        return SUPPORTED_BACKENDS[backend_name_upper]
-
-    # Default to SDPA
-    return SDPABackend

vllm_omni/diffusion/data.py

@@ -14,32 +14,6 @@ from vllm_omni.diffusion.utils.network_utils import is_port_available
 logger = init_logger(__name__)
 
 
-@dataclass
-class TransformerConfig:
-    """Container for raw transformer configuration dictionaries."""
-
-    params: dict[str, Any] = field(default_factory=dict)
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> "TransformerConfig":
-        if not isinstance(data, dict):
-            raise TypeError(f"Expected transformer config dict, got {type(data)!r}")
-        return cls(params=dict(data))
-
-    def to_dict(self) -> dict[str, Any]:
-        return dict(self.params)
-
-    def get(self, key: str, default: Any | None = None) -> Any:
-        return self.params.get(key, default)
-
-    def __getattr__(self, item: str) -> Any:
-        params = object.__getattribute__(self, "params")
-        try:
-            return params[item]
-        except KeyError as exc:
-            raise AttributeError(item) from exc
-
-
 @dataclass
 class OmniDiffusionConfig:
     # Model and path configuration (for convenience)
@@ -47,10 +21,6 @@ class OmniDiffusionConfig:
 
     model_class_name: str | None = None
 
-    dtype: torch.dtype = torch.bfloat16
-
-    tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)
-
     # Attention
     # attention_backend: str = None
 
@@ -110,10 +80,6 @@ class OmniDiffusionConfig:
     vae_cpu_offload: bool = True
     pin_cpu_memory: bool = True
 
-    # VAE memory optimization parameters
-    vae_use_slicing: bool = False
-    vae_use_tiling: bool = False
-
     # STA (Sliding Tile Attention) parameters
     mask_strategy_file_path: str | None = None
     # STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
@@ -242,7 +208,3 @@ class AttentionBackendEnum(enum.Enum):
 
     def __str__(self):
         return self.name.lower()
-
-
-# Special message broadcast via scheduler queues to signal worker shutdown.
-SHUTDOWN_MESSAGE = {"type": "shutdown"}

vllm_omni/diffusion/diffusion_engine.py

@@ -2,15 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import multiprocessing as mp
-import time
 
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, OmniDiffusionConfig
-from vllm_omni.diffusion.registry import get_diffusion_post_process_func, get_diffusion_pre_process_func
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.registry import get_diffusion_post_process_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.scheduler import scheduler
-from vllm_omni.utils.platform_utils import get_diffusion_worker_class
+from vllm_omni.diffusion.worker.gpu_worker import WorkerProc
 
 logger = init_logger(__name__)
 
@@ -27,32 +26,16 @@ class DiffusionEngine:
         self.od_config = od_config
 
         self.post_process_func = get_diffusion_post_process_func(od_config)
-        self.pre_process_func = get_diffusion_pre_process_func(od_config)
 
-        self._processes: list[mp.Process] = []
-        self._closed = False
         self._make_client()
 
     def step(self, requests: list[OmniDiffusionRequest]):
         try:
-            # Apply pre-processing if available
-            if self.pre_process_func is not None:
-                preprocess_start_time = time.time()
-                requests = self.pre_process_func(requests)
-                preprocess_time = time.time() - preprocess_start_time
-                logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds")
-
             output = self.add_req_and_wait_for_response(requests)
             if output.error:
                 raise Exception(f"{output.error}")
             logger.info("Generation completed successfully.")
-
-            postprocess_start_time = time.time()
-            result = self.post_process_func(output.output)
-            postprocess_time = time.time() - postprocess_start_time
-            logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds")
-
-            return result
+            return self.post_process_func(output.output)
         except Exception as e:
             logger.error(f"Generation failed: {e}")
             return None
@@ -85,8 +68,6 @@ class DiffusionEngine:
         else:
             logger.error("Failed to get result queue handle from workers")
 
-        self._processes = processes
-
     def _launch_workers(self, broadcast_handle):
         od_config = self.od_config
         logger.info("Starting server...")
@@ -95,9 +76,6 @@ class DiffusionEngine:
         mp.set_start_method("spawn", force=True)
         processes = []
 
-        # Get the appropriate worker class for current device
-        worker_proc = get_diffusion_worker_class()
-
         # Launch all worker processes
         scheduler_pipe_readers = []
         scheduler_pipe_writers = []
@@ -106,7 +84,7 @@ class DiffusionEngine:
             reader, writer = mp.Pipe(duplex=False)
             scheduler_pipe_writers.append(writer)
             process = mp.Process(
-                target=worker_proc.worker_main,
+                target=WorkerProc.worker_main,
                 args=(
                     i,  # rank
                     od_config,
@@ -150,32 +128,3 @@ class DiffusionEngine:
 
     def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
         return scheduler.add_req(requests)
-
-    def close(self, *, timeout_s: float = 30.0) -> None:
-        if self._closed:
-            return
-        self._closed = True
-
-        # Send shutdown signal to worker processes via broadcast queue
-        try:
-            if getattr(scheduler, "mq", None) is not None:
-                for _ in range(self.od_config.num_gpus or 1):
-                    scheduler.mq.enqueue(SHUTDOWN_MESSAGE)
-        except Exception as exc:  # pragma: no cover - best effort cleanup
-            logger.warning("Failed to send shutdown signal: %s", exc)
-
-        # Join all worker processes, terminate if they refuse to exit
-        for proc in self._processes:
-            if not proc.is_alive():
-                continue
-            proc.join(timeout_s)
-            if proc.is_alive():
-                logger.warning("Terminating diffusion worker %s after timeout", proc.name)
-                proc.terminate()
-                proc.join(timeout_s)
-
-        scheduler.close()
-        self._processes = []
-
-    def __del__(self):  # pragma: no cover - best effort cleanup
-        self.close()

vllm_omni/diffusion/distributed/utils.py

@@ -5,11 +5,7 @@ import os
 
 import torch
 
-from vllm_omni.utils.platform_utils import detect_device_type
-
 
 def get_local_device() -> torch.device:
-    """Return the torch device for the current rank based on detected device type."""
-    device_type = detect_device_type()
-    local_rank = os.environ.get("LOCAL_RANK", 0)
-    return torch.device(f"{device_type}:{local_rank}")
+    """Return the torch device for the current rank."""
+    return torch.device(f"cuda:{os.environ.get('LOCAL_RANK', 0)}")