[Feature] Support online inference (#64)

Signed-off-by: Gaohan123 <gaohan19@huawei.com>

.github/workflows/docs.yml

@@ -1,51 +0,0 @@
-name: Deploy Documentation
-
-on:
-  push:
-    branches:
-      - main
-      - master
-    paths:
-      - 'docs/**'
-      - 'mkdocs.yml'
-      - 'vllm_omni/**'
-      - '.github/workflows/docs.yml'
-  workflow_dispatch:
-
-permissions:
-  contents: write
-  pages: write
-  id-token: write
-
-jobs:
-  deploy:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-          cache: 'pip'
-
-      - name: Install dependencies
-        run: |
-          pip install -e ".[docs]"
-
-      - name: Build documentation
-        run: mkdocs build
-
-      - name: Setup Pages
-        uses: actions/configure-pages@v4
-
-      - name: Upload artifact
-        uses: actions/upload-pages-artifact@v3
-        with:
-          path: ./site
-
-      - name: Deploy to GitHub Pages
-        id: deployment
-        uses: actions/deploy-pages@v4
-

.pre-commit-config.yaml

@@ -26,4 +26,4 @@ repos:
     rev: 7.1.1
     hooks:
       - id: flake8
-        args: ["--max-line-length=88", "--extend-ignore=E203,W503"]
+        args: ["--max-line-length=160", "--extend-ignore=E203,W503"]

README.md

@@ -100,15 +100,7 @@ uv pip install -e .
 
 ## Run examples (Qwen2.5-omni)
 
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen_2_5_omni
-```
-Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run.
-```bash
-bash run.sh
-```
-The output audio is saved in ./output_audio
+Please check the folder of [examples](examples)
 
 ## Further details
 

docs/DOCS_GUIDE.md

@@ -40,15 +40,15 @@ Example docstring:
 ```python
 class OmniLLM:
     """Main entry point for vLLM-omni inference.
-    
+
     This class provides a high-level interface for running multi-modal
     inference with non-autoregressive models.
-    
+
     Args:
         model: Model name or path
         stage_configs: Optional stage configurations
         **kwargs: Additional arguments passed to the engine
-        
+
     Example:
         >>> llm = OmniLLM(model="Qwen/Qwen2.5-Omni")
         >>> outputs = llm.generate(prompts="Hello")

docs/api/overview.md

@@ -66,4 +66,3 @@ Worker classes and model runners for distributed inference.
 - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
-

docs/contributing/design_documents/api_design_doc.md

@@ -2,7 +2,7 @@
 
 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 
+## 📋 Module API
 
 ### Core Module API
 **Core module** provides fundamental scheduling, caching, and resource management functionality.

docs/contributing/design_documents/architecture_overview.md

@@ -1,4 +1,4 @@
-# Architecture Overview 
+# Architecture Overview
 
 # Introduction
 

docs/contributing/design_documents/index.md

@@ -10,4 +10,3 @@ This section contains design documents and architecture specifications for vLLM-
 ## API Design Documentation
 
 - [vLLM-omni API Documentation](api_design_doc.md)
-

docs/index.md

@@ -28,6 +28,5 @@
 
 ## <span class="twemoji">📚</span> Documentation Navigation
 
-- To run open-source models on vLLM-Omni, we recommend starting with the [:material-code-tags: User Quide](user_guide/getting_started/quickstart.md) 
+- To run open-source models on vLLM-Omni, we recommend starting with the [:material-code-tags: User Quide](user_guide/getting_started/quickstart.md)
 - To develop and contribute to vLLM-Omni, we recommend starting with the [:material-tools: Developer Guide](contributing/README.md)
-

docs/mkdocs/hooks/generate_examples.py

@@ -103,11 +103,20 @@ class Example:
         if self.path.is_file():
             return []
         # Text file extensions that can be safely included
-        text_extensions = {".py", ".md", ".sh", ".yaml", ".yml", ".json", ".txt", ".toml", ".cfg", ".ini"}
-        is_other_file = lambda file: (
-            file.is_file() 
-            and file != self.main_file 
-            and file.suffix in text_extensions
+        text_extensions = {
+            ".py",
+            ".md",
+            ".sh",
+            ".yaml",
+            ".yml",
+            ".json",
+            ".txt",
+            ".toml",
+            ".cfg",
+            ".ini",
+        }
+        is_other_file = lambda file: (  # noqa: E731
+            file.is_file() and file != self.main_file and file.suffix in text_extensions
         )
         return [file for file in self.path.rglob("*") if is_other_file(file)]
 
@@ -172,7 +181,7 @@ class Example:
                 f"{code_fence}\n"
             )
         else:
-            with open(self.main_file) as f:
+            with open(self.main_file, encoding="utf-8") as f:
                 # Skip the title from md snippets as it's been included above
                 main_content = f.readlines()[1:]
             content += self.fix_relative_links("".join(main_content))

docs/mkdocs/stylesheets/extra.css

@@ -53,4 +53,3 @@ body[data-md-color-scheme="slate"] .md-nav__item--section > label.md-nav__link .
 .md-nav__item--section:has([href*="api/vllm_omni/index"]) > .md-nav > .md-nav__list > .md-nav__item--nested > .md-nav > .md-nav__list > .md-nav__item--nested.md-nav__item--active > .md-nav {
     display: block;
 }
-

docs/user_guide/examples/index.md

@@ -6,4 +6,3 @@ vLLM-omni's examples are split into two categories:
 - If you are using vLLM-omni from an HTTP application or client, see the *Online Serving* section.
 
 For detailed example documentation, check the [examples directory](https://github.com/vllm-project/vllm-omni/tree/main/examples) in the repository.
-

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

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

docs/user_guide/examples/online_serving/README.md

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

docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md

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

docs/user_guide/getting_started/quickstart.md

@@ -53,4 +53,3 @@ The output audio is saved in ./output_audio
 - Read the [architecture documentation](../../contributing/design_documents/vllm_omni_design.md)
 - Check out the [API reference](../../api/overview.md)
 - Explore the [examples](../examples/index.md)
-

examples/offline_inference/qwen2_5_omni/README.md

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

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -7,8 +7,8 @@ import numpy as np
 import soundfile as sf
 import torch
 from utils import make_omni_prompt
-from vllm.sampling_params import SamplingParams
 
+from vllm.sampling_params import SamplingParams
 from vllm_omni.entrypoints.omni_llm import OmniLLM
 
 _os_env_toggle.environ["VLLM_USE_V1"] = "1"
@@ -109,10 +109,14 @@ def parse_args():
     parser.add_argument("--use-torchvision", action="store_true")
     parser.add_argument("--tokenize", action="store_true")
     parser.add_argument(
-        "--output-wav", default="output.wav", help="[Deprecated] Output wav directory (use --output-dir)."
+        "--output-wav",
+        default="output.wav",
+        help="[Deprecated] Output wav directory (use --output-dir).",
     )
     parser.add_argument(
-        "--output-dir", default="outputs", help="Output directory to save text and wav files together."
+        "--output-dir",
+        default="outputs",
+        help="Output directory to save text and wav files together.",
     )
     parser.add_argument(
         "--thinker-hidden-states-dir",
@@ -168,7 +172,9 @@ def main():
         raise
 
     if args.prompts is None:
-        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
+        raise ValueError(
+            "No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)"
+        )
     omni_llm = OmniLLM(
         model=model_name,
         log_stats=args.enable_stats,
@@ -217,7 +223,9 @@ def main():
     omni_outputs = omni_llm.generate(prompt, sampling_params_list)
 
     # Determine output directory: prefer --output-dir; fallback to --output-wav
-    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    output_dir = (
+        args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    )
     os.makedirs(output_dir, exist_ok=True)
     for stage_outputs in omni_outputs:
         if stage_outputs.final_output_type == "text":

examples/offline_inference/qwen2_5_omni/extract_prompts.py

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

examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh

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

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

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

examples/offline_inference/qwen2_5_omni/utils.py

@@ -9,6 +9,7 @@ import soundfile as sf
 import torch
 import torchvision.io
 from processing_omni import fetch_image, fetch_video
+from transformers import AutoConfig, AutoProcessor
 
 from vllm.inputs import TextPrompt
 from vllm_omni.inputs.data import OmniTokensPrompt

examples/offline_inference/qwen_2_5_omni/README.md

@@ -1,92 +0,0 @@
-# Offline Example of vLLM-omni for Qwen2.5-omni
-
-## 🛠️ Installation
-
-### Environment setup
-
-Use Docker to keep consistent basic environment (Optional, Recommended)
-```bash
-docker run --gpus all --ipc=host --network=host -v $source_dir:$container_dir --name $container_name -it nvcr.io/nvidia/pytorch:25.01-py3 bash
-```
-
-Set up basic uv environment
-```bash
-pip install uv
-uv venv --python 3.12 --seed
-source .venv/bin/activate
-```
-
-### Installation of vLLM (for users)
-
-Now we build it based on vLLM v0.11.0. Please install it with command below.
-```bash
-uv pip install vllm==0.11.0 --torch-backend=auto
-```
-
-### Installation of vLLM (for developers)
-
-If you want to check or debug with source code of vLLM, install stable release version of vllm with 0.11.0 from source with pre-built wheel file.
-
-```bash
-git clone https://github.com/vllm-project/vllm.git
-cd vllm
-git checkout v0.11.0
-```
-Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set VLLM_PRECOMPILED_WHEEL_LOCATION as your local absolute path of whl file.
-```bash
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.11.0/vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl
-```
-Install vllm with command below (If you have no existing PyTorch).
-```bash
-uv pip install --editable .
-```
-
-Install vllm with command below (If you already have PyTorch).
-```bash
-python use_existing_torch.py
-uv pip install -r requirements/build.txt
-uv pip install --no-build-isolation --editable .
-```
-
-### Verification for successful installation of vLLM
-Just run the command below. If no error, it demonstrates that the installation is successfull.
-```bash
-python -c "import vllm._C"
-```
-
-### Installation of vLLM-omni
-Install additional requirements for vllm-omni
-```bash
-git clone https://github.com/vllm-project/vllm-omni.git
-cd vllm_omni
-uv pip install -e .
-```
-
-## Run examples (Qwen2.5-omni)
-### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
-```bash
-tar -xf <Your Download Path>/seedtts_testset.tar
-cp seedtts_testset/en/meta.lst examples/offline_inference/qwen_2_5_omni/meta.lst
-python3 examples/offline_inference/qwen_2_5_omni/extract_prompts.py \
-  --input examples/offline_inference/qwen_2_5_omni/meta.lst \
-  --output examples/offline_inference/qwen_2_5_omni/top100.txt \
-  --topk 100
-```
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen_2_5_omni
-```
-Modify PYTHONPATH in run_multiple_prompts.sh as your path of vllm_omni. Then run.
-```bash
-bash run_multiple_prompts.sh
-```
-### Single Prompts
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen_2_5_omni
-```
-Modify PYTHONPATH in run_single_prompt.sh as your path of vllm_omni. Then run.
-```bash
-bash run_single_prompt.sh
-```

examples/online_serving/README.md

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

examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py

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

examples/online_serving/run_curl_multimodal_generation.sh

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

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vllm-omni"
-version = "0.1.0"
+version = "0.11.0rc1"
 description = "vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures"
 readme = "README.md"
 requires-python = ">=3.9,<3.14"  # Align with vLLM v0.11 supported Python versions

vllm_omni/__init__.py

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

vllm_omni/distributed/__init__.py

@@ -1,2 +0,0 @@
-
-

vllm_omni/engine/arg_utils.py

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

vllm_omni/engine/processor.py

@@ -7,6 +7,7 @@ import torch
 from vllm.config import VllmConfig
 from vllm.inputs import ProcessorInputs, PromptType
 from vllm.inputs.parse import split_enc_dec_inputs
+from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
 from vllm.multimodal.inputs import MultiModalFeatureSpec
@@ -25,6 +26,8 @@ from vllm_omni.engine import (
 )
 from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 
+logger = init_logger(__name__)
+
 
 class OmniProcessor(Processor):
     @staticmethod

vllm_omni/entrypoints/async_omni_llm.py

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

vllm_omni/entrypoints/cli/main.py

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

vllm_omni/entrypoints/cli/serve.py

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

vllm_omni/entrypoints/log_utils.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional
-
 import logging
 import os
 import time
+from typing import Any, Dict, List, Optional
 
 from vllm_omni.entrypoints.stage_utils import append_jsonl as _append_jsonl
 
@@ -49,22 +48,32 @@ def remove_old_logs(log_file: Optional[str], num_stages: int) -> None:
         pass
 
 
-def configure_orchestrator_logger(logger: logging.Logger, log_file: Optional[str]) -> None:
+def configure_orchestrator_logger(
+    logger: logging.Logger, log_file: Optional[str]
+) -> None:
     try:
         if not log_file:
             return
-        has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers)
+        has_file_handler = any(
+            isinstance(h, logging.FileHandler) for h in logger.handlers
+        )
         if not has_file_handler:
             fh = logging.FileHandler(log_file)
             fh.setLevel(logging.DEBUG)
-            fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s"))
+            fh.setFormatter(
+                logging.Formatter(
+                    "%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s"
+                )
+            )
             logger.addHandler(fh)
             logger.setLevel(logging.DEBUG)
     except Exception:
         pass
 
 
-def init_stats_paths(enable_stats: bool, log_file: Optional[str]) -> tuple[Optional[str], Optional[str]]:
+def init_stats_paths(
+    enable_stats: bool, log_file: Optional[str]
+) -> tuple[Optional[str], Optional[str]]:
     stats_file: Optional[str] = None
     overall_file: Optional[str] = None
     try:
@@ -86,122 +95,217 @@ def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None:
         pass
 
 
-def log_transfer_tx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
-                    size_bytes: int, tx_time_ms: float, used_shm: bool) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "transfer_stats",
-        "from_stage": from_stage,
-        "to_stage": to_stage,
-        "request_id": request_id,
-        "size_bytes": int(size_bytes),
-        "tx_time_ms": float(tx_time_ms),
-        "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
-        "used_shm": bool(used_shm),
-    })
+def log_transfer_tx(
+    stats_file: Optional[str],
+    from_stage: int,
+    to_stage: int,
+    request_id: Any,
+    size_bytes: int,
+    tx_time_ms: float,
+    used_shm: bool,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "transfer_stats",
+            "from_stage": from_stage,
+            "to_stage": to_stage,
+            "request_id": request_id,
+            "size_bytes": int(size_bytes),
+            "tx_time_ms": float(tx_time_ms),
+            "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
+            "used_shm": bool(used_shm),
+        },
+    )
 
 
-def log_transfer_rx(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
-                    rx_bytes: int, rx_decode_time_ms: float, in_flight_time_ms: float) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "transfer_rx_stats",
-        "from_stage": from_stage,
-        "to_stage": to_stage,
-        "request_id": request_id,
-        "rx_bytes": int(rx_bytes),
-        "rx_decode_time_ms": float(rx_decode_time_ms),
-        "in_flight_time_ms": float(in_flight_time_ms),
-        "rx_time_per_kb_ms": (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0,
-    })
+def log_transfer_rx(
+    stats_file: Optional[str],
+    from_stage: int,
+    to_stage: int,
+    request_id: Any,
+    rx_bytes: int,
+    rx_decode_time_ms: float,
+    in_flight_time_ms: float,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "transfer_rx_stats",
+            "from_stage": from_stage,
+            "to_stage": to_stage,
+            "request_id": request_id,
+            "rx_bytes": int(rx_bytes),
+            "rx_decode_time_ms": float(rx_decode_time_ms),
+            "in_flight_time_ms": float(in_flight_time_ms),
+            "rx_time_per_kb_ms": (
+                (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6))
+                if rx_bytes > 0
+                else 0.0
+            ),
+        },
+    )
 
 
-def log_transfer_total(stats_file: Optional[str], from_stage: int, to_stage: int, request_id: Any,
-                       size_bytes: int, tx_time_ms: float, in_flight_time_ms: float,
-                       rx_decode_time_ms: float, total_time_ms: float) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "transfer_total_stats",
-        "from_stage": from_stage,
-        "to_stage": to_stage,
-        "request_id": request_id,
-        "size_bytes": int(size_bytes),
-        "tx_time_ms": float(tx_time_ms),
-        "in_flight_time_ms": float(in_flight_time_ms),
-        "rx_decode_time_ms": float(rx_decode_time_ms),
-        "total_time_ms": float(total_time_ms),
-        "total_time_per_kb_ms": float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0,
-    })
+def log_transfer_total(
+    stats_file: Optional[str],
+    from_stage: int,
+    to_stage: int,
+    request_id: Any,
+    size_bytes: int,
+    tx_time_ms: float,
+    in_flight_time_ms: float,
+    rx_decode_time_ms: float,
+    total_time_ms: float,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "transfer_total_stats",
+            "from_stage": from_stage,
+            "to_stage": to_stage,
+            "request_id": request_id,
+            "size_bytes": int(size_bytes),
+            "tx_time_ms": float(tx_time_ms),
+            "in_flight_time_ms": float(in_flight_time_ms),
+            "rx_decode_time_ms": float(rx_decode_time_ms),
+            "total_time_ms": float(total_time_ms),
+            "total_time_per_kb_ms": (
+                float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6)
+                if size_bytes > 0
+                else 0.0
+            ),
+        },
+    )
 
 
-def log_orchestrator_e2e(stats_file: Optional[str], request_id: Any, final_stage_id: int,
-                         e2e_time_ms: float, num_tokens_out: int, e2e_time_per_token_ms: float) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "orchestrator_request_e2e",
-        "request_id": request_id,
-        "final_stage_id": final_stage_id,
-        "e2e_time_ms": float(e2e_time_ms),
-        "num_tokens_out": int(num_tokens_out),
-        "e2e_time_per_token_ms": float(e2e_time_per_token_ms),
-    })
+def log_orchestrator_e2e(
+    stats_file: Optional[str],
+    request_id: Any,
+    final_stage_id: int,
+    e2e_time_ms: float,
+    num_tokens_out: int,
+    e2e_time_per_token_ms: float,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "orchestrator_request_e2e",
+            "request_id": request_id,
+            "final_stage_id": final_stage_id,
+            "e2e_time_ms": float(e2e_time_ms),
+            "num_tokens_out": int(num_tokens_out),
+            "e2e_time_per_token_ms": float(e2e_time_per_token_ms),
+        },
+    )
 
 
-def log_orchestrator_summary(stats_file: Optional[str], summary: Dict[str, Any]) -> None:
+def log_orchestrator_summary(
+    stats_file: Optional[str], summary: Dict[str, Any]
+) -> None:
     _safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary})
 
 
-def log_overall_summary(overall_stats_file: Optional[str], summary: Dict[str, Any]) -> None:
+def log_overall_summary(
+    overall_stats_file: Optional[str], summary: Dict[str, Any]
+) -> None:
     _safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary})
 
 
-def log_overall_record(overall_stats_file: Optional[str], record: Dict[str, Any]) -> None:
+def log_overall_record(
+    overall_stats_file: Optional[str], record: Dict[str, Any]
+) -> None:
     _safe_append_jsonl(overall_stats_file, record)
 
 
-def log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any, batch_size: int,
-                            num_tokens_out: int, stage_gen_time_ms: float, tokens_per_s: float,
-                            rx_transfer_bytes: int, rx_decode_time_ms: float, rx_mbps: float) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "stage_request_stats",
-        "stage_id": stage_id,
-        "request_id": request_id,
-        "batch_size": int(batch_size),
-        "num_tokens_out": int(num_tokens_out),
-        "stage_gen_time_ms": float(stage_gen_time_ms),
-        "tokens_per_s": float(tokens_per_s),
-        "rx_transfer_bytes": int(rx_transfer_bytes),
-        "rx_decode_time_ms": float(rx_decode_time_ms),
-        "rx_mbps": float(rx_mbps),
-    })
+def log_stage_request_stats(
+    stats_file: Optional[str],
+    stage_id: int,
+    request_id: Any,
+    batch_size: int,
+    num_tokens_out: int,
+    stage_gen_time_ms: float,
+    tokens_per_s: float,
+    rx_transfer_bytes: int,
+    rx_decode_time_ms: float,
+    rx_mbps: float,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "stage_request_stats",
+            "stage_id": stage_id,
+            "request_id": request_id,
+            "batch_size": int(batch_size),
+            "num_tokens_out": int(num_tokens_out),
+            "stage_gen_time_ms": float(stage_gen_time_ms),
+            "tokens_per_s": float(tokens_per_s),
+            "rx_transfer_bytes": int(rx_transfer_bytes),
+            "rx_decode_time_ms": float(rx_decode_time_ms),
+            "rx_mbps": float(rx_mbps),
+        },
+    )
 
 
-def log_stage_running_avg(stats_file: Optional[str], stage_id: int, total_tokens: int,
-                          total_gen_time_ms: float, avg_tokens_per_s: float) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "stage_running_avg",
-        "stage_id": stage_id,
-        "total_tokens": int(total_tokens),
-        "total_gen_time_ms": float(total_gen_time_ms),
-        "avg_tokens_per_s": float(avg_tokens_per_s),
-    })
+def log_stage_running_avg(
+    stats_file: Optional[str],
+    stage_id: int,
+    total_tokens: int,
+    total_gen_time_ms: float,
+    avg_tokens_per_s: float,
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "stage_running_avg",
+            "stage_id": stage_id,
+            "total_tokens": int(total_tokens),
+            "total_gen_time_ms": float(total_gen_time_ms),
+            "avg_tokens_per_s": float(avg_tokens_per_s),
+        },
+    )
 
 
-def log_stage_batch_stats(stats_file: Optional[str], stage_id: int, batch_size: int,
-                          batch_gen_time_ms: float, request_ids: List[Any]) -> None:
-    _safe_append_jsonl(stats_file, {
-        "type": "stage_batch_stats",
-        "stage_id": stage_id,
-        "batch_size": int(batch_size),
-        "batch_gen_time_ms": float(batch_gen_time_ms),
-        "request_ids": list(request_ids),
-    })
+def log_stage_batch_stats(
+    stats_file: Optional[str],
+    stage_id: int,
+    batch_size: int,
+    batch_gen_time_ms: float,
+    request_ids: List[Any],
+) -> None:
+    _safe_append_jsonl(
+        stats_file,
+        {
+            "type": "stage_batch_stats",
+            "stage_id": stage_id,
+            "batch_size": int(batch_size),
+            "batch_gen_time_ms": float(batch_gen_time_ms),
+            "request_ids": list(request_ids),
+        },
+    )
 
 
-def compute_and_log_stage_request_stats(stats_file: Optional[str], stage_id: int, request_id: Any,
-                                       batch_size: int, engine_outputs: List[Any],
-                                       stage_gen_time_ms: float, rx_transfer_bytes: int,
-                                       rx_decode_time_ms: float) -> None:
+def compute_and_log_stage_request_stats(
+    stats_file: Optional[str],
+    stage_id: int,
+    request_id: Any,
+    batch_size: int,
+    engine_outputs: List[Any],
+    stage_gen_time_ms: float,
+    rx_transfer_bytes: int,
+    rx_decode_time_ms: float,
+) -> None:
     """Compute per-request metrics and log them in one call."""
     num_tokens = count_tokens_from_outputs(engine_outputs)
-    tokens_per_s = (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
-    rx_mbps = (float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0) if rx_transfer_bytes > 0 else 0.0
+    tokens_per_s = (
+        (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
+    )
+    rx_mbps = (
+        (float(rx_transfer_bytes) * 8.0)
+        / (max(float(rx_decode_time_ms), 1e-6) * 1000.0)
+        if rx_transfer_bytes > 0
+        else 0.0
+    )
     log_stage_request_stats(
         stats_file,
         stage_id,
@@ -218,18 +322,23 @@ def compute_and_log_stage_request_stats(stats_file: Optional[str], stage_id: int
 
 # ----------------- Aggregation helpers for orchestrator -----------------
 
-def record_stage_metrics(per_request: Dict[int, Dict[str, Any]],
-                         stage_req_counts: List[int],
-                         stage_total_time_ms: List[float],
-                         stage_total_tokens: List[int],
-                         stage_id: int,
-                         req_id: Any,
-                         metrics: Dict[str, Any]) -> None:
+
+def record_stage_metrics(
+    per_request: Dict[int, Dict[str, Any]],
+    stage_req_counts: List[int],
+    stage_total_time_ms: List[float],
+    stage_total_tokens: List[int],
+    stage_id: int,
+    req_id: Any,
+    metrics: Dict[str, Any],
+) -> None:
     try:
         stage_req_counts[stage_id] += 1
         stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
         rid_int = int(req_id)
-        pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+        pr = per_request.setdefault(
+            rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}
+        )
         pr_stages = pr["stages"]  # type: ignore[index]
         pr_stages[stage_id] = {
             "stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)),
@@ -239,23 +348,32 @@ def record_stage_metrics(per_request: Dict[int, Dict[str, Any]],
         pass
 
 
-def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
-                                 transfer_agg: Dict[tuple[int, int], Dict[str, float]],
-                                 per_request: Dict[int, Dict[str, Any]],
-                                 stage_id: int,
-                                 req_id: Any,
-                                 rx_bytes: float,
-                                 rx_ms: float,
-                                 in_flight_ms: float) -> Optional[tuple[int, float, float]]:
+def aggregate_rx_and_maybe_total(
+    transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
+    transfer_agg: Dict[tuple[int, int], Dict[str, float]],
+    per_request: Dict[int, Dict[str, Any]],
+    stage_id: int,
+    req_id: Any,
+    rx_bytes: float,
+    rx_ms: float,
+    in_flight_ms: float,
+) -> Optional[tuple[int, float, float]]:
     try:
         # Update RX aggregates for (stage_id-1 -> stage_id)
         if stage_id > 0:
             key = (stage_id - 1, stage_id)
             agg = transfer_agg.get(key)
             if agg is None:
-                agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
-                       "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
-                       "sum_total_ms": 0.0, "total_count": 0.0}
+                agg = {
+                    "sum_bytes": 0.0,
+                    "sum_ms": 0.0,
+                    "count": 0.0,
+                    "sum_rx_bytes": 0.0,
+                    "sum_rx_ms": 0.0,
+                    "rx_count": 0.0,
+                    "sum_total_ms": 0.0,
+                    "total_count": 0.0,
+                }
                 transfer_agg[key] = agg
             agg["sum_rx_bytes"] += float(rx_bytes)
             agg["sum_rx_ms"] += float(rx_ms)
@@ -276,7 +394,9 @@ def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], D
             agg["total_count"] += 1.0
             # accumulate per-request transfer totals
             try:
-                pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+                pr = per_request.setdefault(
+                    rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}
+                )
                 pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms  # type: ignore[index]
                 pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes)  # type: ignore[index]
             except Exception:
@@ -287,20 +407,29 @@ def aggregate_rx_and_maybe_total(transfer_edge_req: Dict[tuple[int, int, int], D
         return None
 
 
-def record_sender_transfer_agg(transfer_agg: Dict[tuple[int, int], Dict[str, float]],
-                               transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
-                               from_stage: int,
-                               to_stage: int,
-                               req_id: Any,
-                               size_bytes: int,
-                               tx_ms: float) -> None:
+def record_sender_transfer_agg(
+    transfer_agg: Dict[tuple[int, int], Dict[str, float]],
+    transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
+    from_stage: int,
+    to_stage: int,
+    req_id: Any,
+    size_bytes: int,
+    tx_ms: float,
+) -> None:
     try:
         key = (from_stage, to_stage)
         agg = transfer_agg.get(key)
         if agg is None:
-            agg = {"sum_bytes": 0.0, "sum_ms": 0.0, "count": 0.0,
-                   "sum_rx_bytes": 0.0, "sum_rx_ms": 0.0, "rx_count": 0.0,
-                   "sum_total_ms": 0.0, "total_count": 0.0}
+            agg = {
+                "sum_bytes": 0.0,
+                "sum_ms": 0.0,
+                "count": 0.0,
+                "sum_rx_bytes": 0.0,
+                "sum_rx_ms": 0.0,
+                "rx_count": 0.0,
+                "sum_total_ms": 0.0,
+                "total_count": 0.0,
+            }
             transfer_agg[key] = agg
         agg["sum_bytes"] += float(size_bytes)
         agg["sum_ms"] += float(tx_ms)
@@ -332,9 +461,11 @@ def count_tokens_from_outputs(engine_outputs: List[Any]) -> int:
     return total
 
 
-def build_stage_summary(stage_req_counts: List[int],
-                        stage_total_tokens: List[int],
-                        stage_total_time_ms: List[float]) -> List[Dict[str, Any]]:
+def build_stage_summary(
+    stage_req_counts: List[int],
+    stage_total_tokens: List[int],
+    stage_total_time_ms: List[float],
+) -> List[Dict[str, Any]]:
     summary: List[Dict[str, Any]] = []
     for sid in range(len(stage_req_counts)):
         reqs = stage_req_counts[sid]
@@ -342,53 +473,74 @@ def build_stage_summary(stage_req_counts: List[int],
         total_ms = float(stage_total_time_ms[sid])
         avg_req = (total_ms / reqs) if reqs > 0 else 0.0
         avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
-        summary.append({
-            "stage_id": sid,
-            "requests": int(reqs),
-            "tokens": int(tokens),
-            "total_time_ms": total_ms,
-            "avg_time_per_request_ms": avg_req,
-            "avg_tokens_per_s": avg_tok,
-        })
+        summary.append(
+            {
+                "stage_id": sid,
+                "requests": int(reqs),
+                "tokens": int(tokens),
+                "total_time_ms": total_ms,
+                "avg_time_per_request_ms": avg_req,
+                "avg_tokens_per_s": avg_tok,
+            }
+        )
     return summary
 
 
-def build_transfer_summary(transfer_agg: Dict[tuple[int, int], Dict[str, float]]) -> List[Dict[str, Any]]:
+def build_transfer_summary(
+    transfer_agg: Dict[tuple[int, int], Dict[str, float]]
+) -> List[Dict[str, Any]]:
     summary: List[Dict[str, Any]] = []
     for (src, dst), agg in transfer_agg.items():
         sum_bytes = float(agg.get("sum_bytes", 0.0))
         sum_ms = float(agg.get("sum_ms", 0.0))
         samples = int(agg.get("count", 0.0))
-        tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
+        tx_mbps = (
+            (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
+        )
         sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0))
         sum_rx_ms = float(agg.get("sum_rx_ms", 0.0))
         samples_rx = int(agg.get("rx_count", 0.0))
-        rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0
+        rx_mbps = (
+            (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0)
+            if sum_rx_bytes > 0
+            else 0.0
+        )
         sum_total_ms = float(agg.get("sum_total_ms", 0.0))
         samples_total = int(agg.get("total_count", 0.0))
-        total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
-        summary.append({
-            "from_stage": src,
-            "to_stage": dst,
-            "samples": samples,
-            "total_bytes": int(sum_bytes),
-            "total_time_ms": sum_ms,
-            "tx_mbps": tx_mbps,
-            "rx_samples": samples_rx,
-            "rx_total_bytes": int(sum_rx_bytes),
-            "rx_total_time_ms": sum_rx_ms,
-            "rx_mbps": rx_mbps,
-            "total_samples": samples_total,
-            "total_transfer_time_ms": sum_total_ms,
-            "total_mbps": total_mbps,
-        })
+        total_mbps = (
+            (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0)
+            if sum_bytes > 0
+            else 0.0
+        )
+        summary.append(
+            {
+                "from_stage": src,
+                "to_stage": dst,
+                "samples": samples,
+                "total_bytes": int(sum_bytes),
+                "total_time_ms": sum_ms,
+                "tx_mbps": tx_mbps,
+                "rx_samples": samples_rx,
+                "rx_total_bytes": int(sum_rx_bytes),
+                "rx_total_time_ms": sum_rx_ms,
+                "rx_mbps": rx_mbps,
+                "total_samples": samples_total,
+                "total_transfer_time_ms": sum_total_ms,
+                "total_mbps": total_mbps,
+            }
+        )
     return summary
 
 
 class OrchestratorMetrics:
-    def __init__(self, num_stages: int, enable_stats: bool,
-                 stats_file: Optional[str], overall_stats_file: Optional[str],
-                 wall_start_ts: float) -> None:
+    def __init__(
+        self,
+        num_stages: int,
+        enable_stats: bool,
+        stats_file: Optional[str],
+        overall_stats_file: Optional[str],
+        wall_start_ts: float,
+    ) -> None:
         self.num_stages = int(num_stages)
         self.enable_stats = bool(enable_stats)
         self.stats_file = stats_file
@@ -406,11 +558,19 @@ class OrchestratorMetrics:
         self.sum_per_request_transfer_ms: float = 0.0
         self.wall_start_ts: float = float(wall_start_ts)
         self.last_finish_ts: float = float(wall_start_ts)
-        self.stage_seen_batches: Dict[int, set] = {sid: set() for sid in range(self.num_stages)}
-        self.stage_first_ts: List[Optional[float]] = [None for _ in range(self.num_stages)]
-        self.stage_last_ts: List[Optional[float]] = [None for _ in range(self.num_stages)]
+        self.stage_seen_batches: Dict[int, set] = {
+            sid: set() for sid in range(self.num_stages)
+        }
+        self.stage_first_ts: List[Optional[float]] = [
+            None for _ in range(self.num_stages)
+        ]
+        self.stage_last_ts: List[Optional[float]] = [
+            None for _ in range(self.num_stages)
+        ]
 
-    def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: Dict[str, Any]) -> None:
+    def on_stage_metrics(
+        self, stage_id: int, req_id: Any, metrics: Dict[str, Any]
+    ) -> None:
         record_stage_metrics(
             self.per_request,
             self.stage_req_counts,
@@ -425,7 +585,9 @@ class OrchestratorMetrics:
             if batch_id_raw is not None:
                 batch_id = int(batch_id_raw)
                 if batch_id not in self.stage_seen_batches[stage_id]:
-                    self.stage_total_time_ms[stage_id] += float(metrics.get("stage_gen_time_ms", 0.0))
+                    self.stage_total_time_ms[stage_id] += float(
+                        metrics.get("stage_gen_time_ms", 0.0)
+                    )
                     self.stage_seen_batches[stage_id].add(batch_id)
         except Exception:
             pass
@@ -466,8 +628,15 @@ class OrchestratorMetrics:
                     float(total_ms_c),
                 )
 
-    def on_forward(self, from_stage: int, to_stage: int, req_id: Any,
-                   size_bytes: int, tx_ms: float, used_shm: bool) -> None:
+    def on_forward(
+        self,
+        from_stage: int,
+        to_stage: int,
+        req_id: Any,
+        size_bytes: int,
+        tx_ms: float,
+        used_shm: bool,
+    ) -> None:
         # Mark first input time for the destination stage if not set
         if self.stage_first_ts[to_stage] is None:
             self.stage_first_ts[to_stage] = time.time()
@@ -491,8 +660,9 @@ class OrchestratorMetrics:
             float(tx_ms),
         )
 
-    def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: List[Any],
-                            req_start_ts: float) -> None:
+    def on_finalize_request(
+        self, stage_id: int, req_id: Any, engine_outputs: List[Any], req_start_ts: float
+    ) -> None:
         try:
             rid_int = int(req_id)
         except Exception:
@@ -509,7 +679,9 @@ class OrchestratorMetrics:
         self.e2e_total_tokens += int(num_tokens)
         self.e2e_count += 1
         self.e2e_done.add(rid_int)
-        pr = self.per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+        pr = self.per_request.setdefault(
+            rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}
+        )
         per_req_record = {
             "type": "overall_request",
             "request_id": rid_int,
@@ -524,7 +696,9 @@ class OrchestratorMetrics:
             log_overall_record(self.overall_stats_file, per_req_record)
         if self.enable_stats and self.stats_file:
             e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0
-            log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt)
+            log_orchestrator_e2e(
+                self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt
+            )
 
     def build_and_log_summary(self, final_stage_id_for_e2e: int) -> Dict[str, Any]:
         # Compute stage summary using wall time between first input and last output per stage
@@ -532,22 +706,34 @@ class OrchestratorMetrics:
         for sid in range(self.num_stages):
             first_ts = self.stage_first_ts[sid]
             last_ts = self.stage_last_ts[sid]
-            total_ms = (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0
+            total_ms = (
+                (max(0.0, (last_ts - first_ts)) * 1000.0)
+                if (first_ts is not None and last_ts is not None)
+                else 0.0
+            )
             reqs = self.stage_req_counts[sid]
             tokens = self.stage_total_tokens[sid]
             avg_req = (total_ms / reqs) if reqs > 0 else 0.0
             avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
-            stage_summary.append({
-                "stage_id": sid,
-                "requests": int(reqs),
-                "tokens": int(tokens),
-                "total_time_ms": float(total_ms),
-                "avg_time_per_request_ms": float(avg_req),
-                "avg_tokens_per_s": float(avg_tok),
-            })
+            stage_summary.append(
+                {
+                    "stage_id": sid,
+                    "requests": int(reqs),
+                    "tokens": int(tokens),
+                    "total_time_ms": float(total_ms),
+                    "avg_time_per_request_ms": float(avg_req),
+                    "avg_tokens_per_s": float(avg_tok),
+                }
+            )
         transfer_summary = build_transfer_summary(self.transfer_agg)
-        e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
-        e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0
+        e2e_avg_req = (
+            (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
+        )
+        e2e_avg_tok = (
+            (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms)
+            if self.e2e_total_ms > 0
+            else 0.0
+        )
         wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
         summary: Dict[str, Any] = {
             "e2e_requests": int(self.e2e_count),
@@ -566,5 +752,3 @@ class OrchestratorMetrics:
         if self.enable_stats and self.overall_stats_file:
             log_overall_summary(self.overall_stats_file, summary)
         return summary
-
-

vllm_omni/entrypoints/omni_llm.py

@@ -1,11 +1,9 @@
-from typing import Any, Optional, Sequence, Union
-import logging
 import multiprocessing as mp
-from concurrent.futures import ThreadPoolExecutor, as_completed
 import os
-import queue
 import sys
 import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Optional, Sequence, Union
 
 import cloudpickle
 from pydantic import ValidationError
@@ -26,18 +24,19 @@ from vllm.v1.engine.llm_engine import LLMEngine
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
-from vllm_omni.entrypoints.omni_stage import OmniStage
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model
 from vllm_omni.entrypoints.log_utils import (
-    remove_old_logs,
+    OrchestratorMetrics,
     configure_orchestrator_logger,
     init_stats_paths,
-    OrchestratorMetrics,
+    remove_old_logs,
 )
-from vllm_omni.entrypoints.stage_utils import (
-    maybe_load_from_ipc as _load,
-    encode_for_ipc as _encode,
-    serialize_obj as _ser,
+from vllm_omni.entrypoints.omni_stage import OmniStage
+from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
+from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
+from vllm_omni.entrypoints.stage_utils import serialize_obj as _ser
+from vllm_omni.entrypoints.utils import (
+    load_stage_configs_from_model,
+    load_stage_configs_from_yaml,
 )
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -52,23 +51,25 @@ class OmniLLM:
     - max_inflight=1 per stage (serial within a stage), but pipeline across stages
     """
 
-    def __init__(self, model: str,
-                 stage_configs=None,
-                 log_stats: bool = False,
-                 log_file: Optional[str] = None,
-                 init_sleep_seconds: int = 20,
-                 shm_threshold_bytes: int = 65536,
-                 batch_timeout: int = 10,
-                 init_timeout: int = 300,
-                 **kwargs):
+    def __init__(
+        self,
+        model: str,
+        stage_configs_path: Optional[str] = None,
+        log_stats: bool = False,
+        log_file: Optional[str] = None,
+        init_sleep_seconds: int = 20,
+        shm_threshold_bytes: int = 65536,
+        batch_timeout: int = 10,
+        init_timeout: int = 300,
+        **kwargs,
+    ):
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
         # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
-        if stage_configs is None:
+        if stage_configs_path is None:
             self.stage_configs = load_stage_configs_from_model(model)
         else:
-            self.stage_configs = stage_configs
-
+            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
 
         # Optional file handler for orchestrator
         self._log_file = log_file
@@ -76,18 +77,34 @@ class OmniLLM:
             remove_old_logs(self._log_file, len(self.stage_list))
             configure_orchestrator_logger(logger, self._log_file)
 
-        self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
-        self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
+        self._stats_file, self._overall_stats_file = init_stats_paths(
+            self._enable_stats, self._log_file
+        )
+        self._initialize_stages(
+            model, init_sleep_seconds, shm_threshold_bytes, init_timeout
+        )
 
-    def _initialize_stages(self, model: str, init_sleep_seconds: int, shm_threshold_bytes: int, init_timeout: int) -> None:
+    def _initialize_stages(
+        self,
+        model: str,
+        init_sleep_seconds: int,
+        shm_threshold_bytes: int,
+        init_timeout: int,
+    ) -> None:
         self.stage_list: list[OmniStage] = []
+
         # Build OmniStage instances in parallel, preserve original order
         def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
             idx, cfg = idx_cfg
             return idx, OmniStage(cfg)
 
-        with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor:
-            futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)]
+        with ThreadPoolExecutor(
+            max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))
+        ) as executor:
+            futures = [
+                executor.submit(_build_stage, (idx, cfg))
+                for idx, cfg in enumerate(self.stage_configs)
+            ]
             results: list[tuple[int, OmniStage]] = []
             for fut in as_completed(futures):
                 results.append(fut.result())
@@ -104,7 +121,6 @@ class OmniLLM:
         # Wait for all stages to report readiness before seeding
         self._stages_ready: set[int] = set()
         self._wait_for_stages_ready(timeout=init_timeout)
-        
 
     def _start_stage_processes(self, model: str) -> None:
         for stage_id, stage in enumerate(self.stage_list):
@@ -131,7 +147,10 @@ class OmniLLM:
             try:
                 q.put_nowait(None)
             except Exception as e:
-                logger.warning("[Orchestrator] Failed to send shutdown signal to stage input queue: %s", e)
+                logger.warning(
+                    "[Orchestrator] Failed to send shutdown signal to stage input queue: %s",
+                    e,
+                )
         for stage in self.stage_list:
             try:
                 stage.stop_stage_worker()
@@ -150,10 +169,27 @@ class OmniLLM:
         sampling_params_list: Optional[
             Union[SamplingParams, Sequence[SamplingParams]]
         ] = None,
+    ) -> list[OmniRequestOutput]:
+        try:
+            return self._run_generation(prompts, sampling_params_list)
+        except Exception as e:
+            logger.exception("[Orchestrator] Failed to run generation: %s", e)
+            raise e
+        finally:
+            self.close()
+
+    def _run_generation(
+        self,
+        prompts: Union[PromptType, Sequence[PromptType]],
+        sampling_params_list: Optional[
+            Union[SamplingParams, Sequence[SamplingParams]]
+        ] = None,
     ) -> list[OmniRequestOutput]:
         logger.debug("[Orchestrator] generate() called")
         if sampling_params_list is None:
-            raise ValueError("sampling_params_list is required for pipelined generation")
+            raise ValueError(
+                "sampling_params_list is required for pipelined generation"
+            )
         if len(sampling_params_list) != len(self.stage_list):
             raise ValueError(
                 f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}"
@@ -171,12 +207,13 @@ class OmniLLM:
         num_stages = len(self.stage_list)
 
         # Map from request_id to original prompt
-        request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
+        request_id_to_prompt: dict[int, PromptType] = {
+            i: p for i, p in enumerate(request_prompts)
+        }
 
         # Track per-request start time for end-to-end timing
         _req_start_ts: dict[int, float] = {}
         _wall_start_ts: float = time.time()
-        _last_finish_ts: float = _wall_start_ts
 
         # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
         final_stage_id_for_e2e = -1
@@ -187,13 +224,25 @@ class OmniLLM:
             if final_stage_id_for_e2e < 0:
                 final_stage_id_for_e2e = len(self.stage_list) - 1
         except Exception as e:
-            logger.debug("[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s", e, exc_info=True)
+            logger.debug(
+                "[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s",
+                e,
+                exc_info=True,
+            )
             final_stage_id_for_e2e = len(self.stage_list) - 1
         # Metrics/aggregation helper
-        metrics = OrchestratorMetrics(num_stages, self._enable_stats, self._stats_file, self._overall_stats_file, _wall_start_ts)
+        metrics = OrchestratorMetrics(
+            num_stages,
+            self._enable_stats,
+            self._stats_file,
+            self._overall_stats_file,
+            _wall_start_ts,
+        )
 
         # Seed stage-0 queue with all requests
-        logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts))
+        logger.debug(
+            "[Orchestrator] Seeding %d requests into stage-0", len(request_prompts)
+        )
         # Mark first input time for stage-0
         metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
 
@@ -214,7 +263,11 @@ class OmniLLM:
         completed_requests = 0
         total_requests = len(request_prompts)
 
-        logger.debug("[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", total_requests, num_stages)
+        logger.debug(
+            "[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d",
+            total_requests,
+            num_stages,
+        )
         while completed_requests < total_requests:
             made_progress = False
             for stage_id, stage in enumerate(self.stage_list):
@@ -225,24 +278,42 @@ class OmniLLM:
                 made_progress = True
                 req_id = result.get("request_id")
                 if "error" in result:
-                    logger.error("Stage %s error on request %s: %s", stage_id, req_id, result["error"])
+                    logger.error(
+                        "Stage %s error on request %s: %s",
+                        stage_id,
+                        req_id,
+                        result["error"],
+                    )
                     continue
-                
+
                 if result.get("type") == "stage_ready":
-                    #Only happens when stage is initialized slower than expected, so we wait for a short time and try again
+                    # Only happens when stage is initialized slower than expected, so we wait for a short time and try again
                     time.sleep(0.05)
                     continue
 
-                engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
+                engine_outputs = _load(
+                    result, obj_key="engine_outputs", shm_key="engine_outputs_shm"
+                )
                 # Mark last output time for this stage whenever we receive outputs
-                metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
+                metrics.stage_last_ts[stage_id] = max(
+                    metrics.stage_last_ts[stage_id] or 0.0, time.time()
+                )
                 try:
                     _m = result.get("metrics")
                     if _m is not None:
                         metrics.on_stage_metrics(stage_id, req_id, _m)
                 except Exception as e:
-                    logger.exception("[Orchestrator] Failed to process metrics for stage %s, req %s: %s", stage_id, req_id, e)
-                logger.debug("[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", stage_id, req_id)
+                    logger.exception(
+                        "[Orchestrator] Failed to process metrics for stage %s, req %s: %s",
+                        stage_id,
+                        req_id,
+                        e,
+                    )
+                logger.debug(
+                    "[Orchestrator] Stage-%s completed request %s; forwarding or finalizing",
+                    stage_id,
+                    req_id,
+                )
                 stage.set_engine_outputs(engine_outputs)
 
                 if getattr(stage, "final_output", False):
@@ -253,20 +324,43 @@ class OmniLLM:
                             request_output=engine_outputs,
                         )
                     )
-                    logger.debug("[Orchestrator] Request %s finalized at stage-%s", req_id, stage_id)
+                    logger.debug(
+                        "[Orchestrator] Request %s finalized at stage-%s",
+                        req_id,
+                        stage_id,
+                    )
 
                     # End-to-end timing and time-per-token for final output (only once per request at the designated final stage)
                     try:
-                        rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
-                        if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
-                            metrics.on_finalize_request(stage_id, req_id, engine_outputs, _req_start_ts.get(req_id, _wall_start_ts))
+                        rid_int = (
+                            int(req_id)
+                            if isinstance(req_id, (int, str)) and str(req_id).isdigit()
+                            else req_id
+                        )
+                        if (
+                            stage_id == final_stage_id_for_e2e
+                            and rid_int not in metrics.e2e_done
+                        ):
+                            metrics.on_finalize_request(
+                                stage_id,
+                                req_id,
+                                engine_outputs,
+                                _req_start_ts.get(req_id, _wall_start_ts),
+                            )
                     except Exception as e:
-                        logger.exception("[Orchestrator] Finalize request handling error for req %s at stage %s: %s", req_id, stage_id, e)
+                        logger.exception(
+                            "[Orchestrator] Finalize request handling error for req %s at stage %s: %s",
+                            req_id,
+                            stage_id,
+                            e,
+                        )
 
                 next_stage_id = stage_id + 1
                 if next_stage_id < num_stages:
                     next_stage: OmniStage = self.stage_list[next_stage_id]
-                    next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
+                    next_inputs = next_stage.process_engine_inputs(
+                        self.stage_list, [request_id_to_prompt[req_id]]
+                    )
                     sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
                     try:
                         # Measure transfer size and time (encode + enqueue)
@@ -282,27 +376,51 @@ class OmniLLM:
                             obj_key="engine_inputs",
                             shm_key="engine_inputs_shm",
                         )
-                        ipc_payload.update({
-                            "request_id": req_id,
-                            "sampling_params": sp_next,
-                            "sent_ts": time.time(),
-                        })
+                        ipc_payload.update(
+                            {
+                                "request_id": req_id,
+                                "sampling_params": sp_next,
+                                "sent_ts": time.time(),
+                            }
+                        )
                         self.stage_list[next_stage_id].submit(ipc_payload)
                         t1 = time.time()
                         tx_ms = (t1 - t0) * 1000.0
-                        metrics.on_forward(stage_id, next_stage_id, req_id, int(size_bytes), float(tx_ms), bool("engine_inputs_shm" in ipc_payload))
+                        metrics.on_forward(
+                            stage_id,
+                            next_stage_id,
+                            req_id,
+                            int(size_bytes),
+                            float(tx_ms),
+                            bool("engine_inputs_shm" in ipc_payload),
+                        )
                     except Exception as e:
-                        logger.warning("[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload", req_id, e)
-                        self.stage_list[next_stage_id].submit({
-                            "request_id": req_id,
-                            "engine_inputs": next_inputs,
-                            "sampling_params": sp_next,
-                        })
-                    logger.debug("[Orchestrator] Forwarded request %s to stage-%s", req_id, next_stage_id)
+                        logger.warning(
+                            "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
+                            req_id,
+                            e,
+                        )
+                        self.stage_list[next_stage_id].submit(
+                            {
+                                "request_id": req_id,
+                                "engine_inputs": next_inputs,
+                                "sampling_params": sp_next,
+                            }
+                        )
+                    logger.debug(
+                        "[Orchestrator] Forwarded request %s to stage-%s",
+                        req_id,
+                        next_stage_id,
+                    )
                     remaining_by_stage[next_stage_id] += 1
                 else:
                     completed_requests += 1
-                    logger.debug("[Orchestrator] Request %s fully completed (%d/%d)", req_id, completed_requests, total_requests)
+                    logger.debug(
+                        "[Orchestrator] Request %s fully completed (%d/%d)",
+                        req_id,
+                        completed_requests,
+                        total_requests,
+                    )
 
             if not made_progress:
                 time.sleep(0.005)
@@ -341,7 +459,9 @@ class OmniLLM:
             not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
             logger.warning(
                 "[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s",
-                len(self._stages_ready), num_stages, not_ready,
+                len(self._stages_ready),
+                num_stages,
+                not_ready,
             )
             # Provide actionable suggestions before shutdown
             try:

vllm_omni/entrypoints/omni_stage.py

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

vllm_omni/entrypoints/openai/api_server.py

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

vllm_omni/entrypoints/openai/serving_chat.py

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

vllm_omni/entrypoints/stage_utils.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Optional, Tuple, Union
-
-import logging
 import json
+import logging
 import os
 import pickle
 from multiprocessing import shared_memory as _shm
+from typing import Any, Dict, Optional, Tuple, Union
 
 import cloudpickle
+from omegaconf import OmegaConf
 
 logger = logging.getLogger(__name__)
 
@@ -37,12 +37,18 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
                     selected_physical = int(toks[0])
                     logger.debug(
                         "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s",
-                        stage_id, devices, selected_physical,
+                        stage_id,
+                        devices,
+                        selected_physical,
                     )
                 except Exception as e:
-                    logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
+                    logger.debug(
+                        "[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e
+                    )
                     selected_physical = None
-        elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()):
+        elif isinstance(devices, (int, str)) and (
+            isinstance(devices, int) or str(devices).isdigit()
+        ):
             logical_idx = max(0, int(devices))
             vis = os.environ.get("CUDA_VISIBLE_DEVICES")
             if vis:
@@ -51,43 +57,72 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
                     if 0 <= logical_idx < len(mapping):
                         selected_physical = mapping[logical_idx]
                 except Exception as e:
-                    logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e)
+                    logger.debug(
+                        "[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s",
+                        stage_id,
+                        e,
+                    )
                     selected_physical = None
             if selected_physical is None:
                 selected_physical = int(logical_idx)
             os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
             logger.debug(
                 "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES to single device",
-                stage_id, logical_idx + 1, selected_physical,
+                stage_id,
+                logical_idx + 1,
+                selected_physical,
             )
         elif devices in (None, "cpu"):
-            logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
+            logger.debug(
+                "[Stage-%s] Using default device visibility (devices=%s)",
+                stage_id,
+                devices,
+            )
         else:
             selected_physical = int(str(devices))
             os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
-            logger.debug("[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical)
+            logger.debug(
+                "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)",
+                stage_id,
+                selected_physical,
+            )
 
         try:
             import torch  # noqa: WPS433
+
             if torch.cuda.is_available():
                 try:
                     torch.cuda.set_device(0)
                 except Exception as e:
-                    logger.debug("[Stage-%s] torch.cuda.set_device(0) failed: %s", stage_id, e, exc_info=True)
+                    logger.debug(
+                        "[Stage-%s] torch.cuda.set_device(0) failed: %s",
+                        stage_id,
+                        e,
+                        exc_info=True,
+                    )
                 num = torch.cuda.device_count()
                 info = []
                 for i in range(num):
                     total = torch.cuda.get_device_properties(i).total_memory
                     free, _ = torch.cuda.mem_get_info(i)
-                    info.append({
-                        "idx": i,
-                        "name": torch.cuda.get_device_name(i),
-                        "total": int(total),
-                        "free": int(free),
-                    })
-                logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info)
+                    info.append(
+                        {
+                            "idx": i,
+                            "name": torch.cuda.get_device_name(i),
+                            "total": int(total),
+                            "free": int(free),
+                        }
+                    )
+                logger.debug(
+                    "[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info
+                )
         except Exception as e:
-            logger.debug("[Stage-%s] Failed to query CUDA devices: %s", stage_id, e, exc_info=True)
+            logger.debug(
+                "[Stage-%s] Failed to query CUDA devices: %s",
+                stage_id,
+                e,
+                exc_info=True,
+            )
     except Exception as e:
         logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)
 
@@ -110,7 +145,7 @@ def shm_write_bytes(payload: bytes) -> Dict[str, Any]:
     try:
         shm.close()
     except Exception as e:
-        logger.debug("Ensure parent dir failed for %s: %s", path, e)
+        logger.debug("Failed to close shared memory: %s", e)
     return meta
 
 
@@ -180,7 +215,9 @@ def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) -
     return container[obj_key]
 
 
-def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, shm_key: str) -> tuple[Any, Dict[str, float]]:
+def maybe_load_from_ipc_with_metrics(
+    container: Dict[str, Any], obj_key: str, shm_key: str
+) -> tuple[Any, Dict[str, float]]:
     """Load object and return (object, metrics) with RX bytes and decode time.
 
     Metrics keys:
@@ -188,6 +225,7 @@ def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, sh
       - rx_decode_time_ms: float
     """
     import time as _time  # local import to avoid overhead at module import
+
     t0 = _time.time()
     if shm_key in container:
         meta = container[shm_key]  # type: ignore[index]
@@ -211,7 +249,9 @@ def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, sh
     }
 
 
-def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> Dict[str, Any]:
+def encode_for_ipc(
+    obj: Any, threshold: int, obj_key: str, shm_key: str
+) -> Dict[str, Any]:
     """Return a dict payload for IPC: inline (obj_key) or SHM (shm_key).
 
     When serialized size exceeds threshold, returns {shm_key: {name,size}};

vllm_omni/entrypoints/utils.py

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

vllm_omni/inputs/preprocess.py

@@ -4,10 +4,13 @@ from typing_extensions import assert_never
 
 from vllm.inputs.data import SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt
 from vllm.inputs.preprocess import InputPreprocessor
+from vllm.logger import init_logger
 from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict
 from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni
 from vllm_omni.inputs.parse import parse_singleton_prompt_omni
 
+logger = init_logger(__name__)
+
 
 class OmniInputPreprocessor(InputPreprocessor):
     def _process_tokens(

vllm_omni/model_executor/models/qwen2_5_omni.py

@@ -279,7 +279,9 @@ class Qwen2_5OmniForConditionalGeneration(
 
             if input_ids is None and additional_information is None:
                 input_ids = torch.zeros(
-                    inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device
+                    inputs_embeds.shape[0],
+                    dtype=torch.long,
+                    device=inputs_embeds.device,
                 )
                 additional_information = {}
                 self.thinker_reply_part = torch.zeros_like(inputs_embeds)
@@ -293,13 +295,19 @@ class Qwen2_5OmniForConditionalGeneration(
 
             # ------- Request-scoped additional information (no cross-request concat) -------
             request_ids: Optional[list[str]] = kwargs.get("request_ids")  # ordered
-            request_token_spans: Optional[list[tuple[int,int]]] = kwargs.get("request_token_spans")
+            request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get(
+                "request_token_spans"
+            )
             addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id")
             runtime_addi = kwargs.get("runtime_additional_information")
 
             # Normalize runtime_addi into a mapping by request_id for convenience
             runtime_addi_by_req: dict[str, dict] = {}
-            if isinstance(request_ids, list) and isinstance(runtime_addi, list) and len(runtime_addi) == len(request_ids):
+            if (
+                isinstance(request_ids, list)
+                and isinstance(runtime_addi, list)
+                and len(runtime_addi) == len(request_ids)
+            ):
                 for i, rid in enumerate(request_ids):
                     if isinstance(rid, str) and isinstance(runtime_addi[i], dict):
                         runtime_addi_by_req[rid] = runtime_addi[i]
@@ -332,9 +340,19 @@ class Qwen2_5OmniForConditionalGeneration(
                     otoks = info.get("thinker_output_token_ids")  # list[int]
 
                     if not isinstance(pe, torch.Tensor):
-                        pe = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
+                        pe = torch.zeros(
+                            0,
+                            self.talker.config.hidden_size,
+                            dtype=inputs_embeds.dtype,
+                            device=self._module_device(self.model),
+                        )
                     if not isinstance(tr, torch.Tensor):
-                        tr = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
+                        tr = torch.zeros(
+                            0,
+                            self.talker.config.hidden_size,
+                            dtype=inputs_embeds.dtype,
+                            device=self._module_device(self.model),
+                        )
                     if not isinstance(ptoks, (list, torch.Tensor)):
                         ptoks = []
                     if not isinstance(otoks, (list, torch.Tensor)):
@@ -342,22 +360,35 @@ class Qwen2_5OmniForConditionalGeneration(
 
                     req_input_ids, req_embeds = self._thinker_to_talker_prefill(
                         voice_type=voice_type,
-                        output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                        output_prompt_embeds=tr.to(inputs_embeds.dtype).to(
+                            self._module_device(self.model)
+                        ),
                         output_token_ids=otoks,
-                        thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                        thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(
+                            self._module_device(self.model)
+                        ),
                         prompt_token_ids=ptoks,
                     )
                     seg_len = min(span_len, req_embeds.shape[0])
                     inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
-                    if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
+                    if (
+                        isinstance(req_input_ids, torch.Tensor)
+                        and req_input_ids.numel() == seg_len
+                    ):
                         input_ids[s : s + seg_len] = req_input_ids
 
                     # Prepare per-request reply queue for subsequent decode: drop first row
                     if tr.ndim == 2 and tr.shape[0] > 0:
-                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = tr[1:].detach().to("cpu").contiguous()
+                        update_by_req_id.setdefault(rid, {})[
+                            "thinker_reply_part_per_request"
+                        ] = (tr[1:].detach().to("cpu").contiguous())
 
             # ------- Decode: span_len == 1 -------
-            if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list):
+            if (
+                not is_profile
+                and isinstance(request_ids, list)
+                and isinstance(request_token_spans, list)
+            ):
                 for idx_req, rid in enumerate(request_ids):
                     s, e = request_token_spans[idx_req]
                     if (int(e) - int(s)) != 1:
@@ -367,18 +398,34 @@ class Qwen2_5OmniForConditionalGeneration(
                     # A) runtime queue
                     q = None
                     if isinstance(rid, str):
-                        q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request")
+                        q = runtime_addi_by_req.get(rid, {}).get(
+                            "thinker_reply_part_per_request"
+                        )
                     if isinstance(q, torch.Tensor) and q.numel() > 0:
                         step_vec = q[0:1]
                         new_q = q[1:].detach().to("cpu").contiguous()
-                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q
+                        update_by_req_id.setdefault(rid, {})[
+                            "thinker_reply_part_per_request"
+                        ] = new_q
                     else:
                         # B) per-request provided decode vector (optional)
-                        info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {}
-                        dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None
+                        info = (
+                            addi_by_req.get(rid, {})
+                            if isinstance(addi_by_req, dict)
+                            else {}
+                        )
+                        dv = (
+                            info.get("decode_output_prompt_embeds")
+                            if isinstance(info, dict)
+                            else None
+                        )
                         if isinstance(dv, torch.Tensor) and dv.numel() > 0:
                             step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
-                        elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0:
+                        elif (
+                            hasattr(self, "thinker_reply_part")
+                            and isinstance(self.thinker_reply_part, torch.Tensor)
+                            and self.thinker_reply_part.numel() > 0
+                        ):
                             # C) fallback shared pool
                             step_vec = self.thinker_reply_part[0:1]
                             self.thinker_reply_part = self.thinker_reply_part[1:]
@@ -386,7 +433,9 @@ class Qwen2_5OmniForConditionalGeneration(
                     if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
                         one_id = input_ids[s : s + 1]
                         _, one_embed = self._thinker_to_talker_decode_one_step(
-                            output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                            output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(
+                                self._module_device(self.model)
+                            ),
                             output_token_ids=one_id,
                         )
                         inputs_embeds[s] = one_embed[0]
@@ -400,11 +449,15 @@ class Qwen2_5OmniForConditionalGeneration(
             multimodal_outputs: dict = None
             # Return updates if any
             if update_by_req_id:
-                multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id}
+                multimodal_outputs = {
+                    "additional_information_update_by_req_id": update_by_req_id
+                }
 
             if sampling_metadata is not None:
                 # the padding token id is set to text model's pad token id, which do not match with the talker model's word embedding size
-                sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448
+                sampling_metadata.prompt_token_ids[
+                    sampling_metadata.prompt_token_ids == 152064
+                ] = 8448
 
             return OmniOutput(
                 text_hidden_states=talker_hidden,
@@ -421,13 +474,11 @@ class Qwen2_5OmniForConditionalGeneration(
                     device=inputs_embeds.device,
                 )
             )
-            
+
             code = code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code
             code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code
-            
-            audio_tensor = self.generate_audio(
-                code, voice_type
-            )
+
+            audio_tensor = self.generate_audio(code, voice_type)
             return OmniOutput(
                 text_hidden_states=None, multimodal_outputs={"audio": audio_tensor}
             )

vllm_omni/model_executor/models/qwen2_5_omni_talker.py

@@ -166,8 +166,16 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         )
         return hidden_states
 
+    def bad_word_processor(self, logits: torch.Tensor) -> torch.Tensor:
+        if hasattr(self.config, "tts_codec_start_token_id"):
+            bos_id = int(getattr(self.config, "tts_codec_start_token_id"))
+            logits[..., bos_id] = -1e9
+        return logits
+
     def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:
-        return self.language_model.compute_logits(hidden_states)
+        logits = self.language_model.compute_logits(hidden_states)
+        logits = self.bad_word_processor(logits)
+        return logits
 
     def sample(
         self,

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -15,8 +15,17 @@ stage_args:
       trust_remote_code: true
       engine_output_type: latent  # change the param name,such as pooling_output
       enable_prefix_caching: false
+    is_comprehension: true
     final_output: true
     final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
   - stage_id: 1
     runtime:
       process: true
@@ -34,6 +43,15 @@ stage_args:
       engine_output_type: latent
     engine_input_source: [0]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
+    default_sampling_params:
+      temperature: 0.9
+      top_p: 0.8
+      top_k: 40
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+      stop_token_ids: [8294]
   - stage_id: 2
     runtime:
       process: true
@@ -52,6 +70,14 @@ stage_args:
     engine_input_source: [1]
     final_output: true
     final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
 
 # Top-level runtime config (concise): default windows and stage edges
 runtime:

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -4,9 +4,12 @@ import torch
 
 from vllm.inputs import TextPrompt
 from vllm_omni.inputs.data import OmniTokensPrompt
+
 TALKER_CODEC_PAD_TOKEN_ID = 8292
 TALKER_CODEC_START_TOKEN_ID = 8293
 TALKER_CODEC_END_TOKEN_ID = 8294
+
+
 def thinker2talker(
     stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None
 ):
@@ -19,9 +22,11 @@ def thinker2talker(
         raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
     thinker_outputs = stage_list[source_stage_id].engine_outputs
     talker_inputs = []
+    if not isinstance(prompt, list):
+        prompt = [prompt]
     multi_modal_data = {
-        thinker_output.request_id: prompt.get("multi_modal_data", None)
-        for thinker_output, prompt in zip(thinker_outputs, prompt)
+        thinker_output.request_id: p.get("multi_modal_data", None)
+        for thinker_output, p in zip(thinker_outputs, prompt)
     }
 
     for i, thinker_output in enumerate(thinker_outputs):
@@ -31,21 +36,28 @@ def thinker2talker(
         prompt_token_ids_len = len(prompt_token_ids)
         latent = output.multimodal_output["latent"]
         thinker_hidden_states = latent.clone().detach().to(latent.device)
+        addtional_information = {
+            "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(
+                torch.float32
+            ),
+            "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(
+                torch.float32
+            ),
+            "prompt_token_ids": prompt_token_ids,
+            "thinker_output_token_ids": thinker_output_ids,
+            "thinker_result_shape": list(
+                thinker_hidden_states[prompt_token_ids_len:].shape
+            ),
+            "prompt_embeds_shape": list(
+                thinker_hidden_states[:prompt_token_ids_len].shape
+            ),
+        }
         talker_inputs.append(
             OmniTokensPrompt(
-                prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID] + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids)) + [TALKER_CODEC_END_TOKEN_ID],
-                additional_information={
-                    "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(
-                        torch.float32
-                    ),
-                    "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(
-                        torch.float32
-                    ),
-                    "prompt_token_ids": prompt_token_ids,
-                    "thinker_output_token_ids": thinker_output_ids,
-                    "thinker_result_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape),
-                    "prompt_embeds_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape),
-                },
+                prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID]
+                + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids))
+                + [TALKER_CODEC_END_TOKEN_ID],
+                additional_information=addtional_information,
                 multi_modal_data=(
                     multi_modal_data[thinker_output.request_id]
                     if multi_modal_data is not None

vllm_omni/version.py

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

vllm_omni/worker/gpu_ar_model_runner.py

@@ -138,9 +138,15 @@ class GPUARModelRunner(OmniGPUModelRunner):
                             # For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor
                             if overlay_len > 0:
                                 try:
-                                    seg = v[
-                                        num_computed_tokens : num_computed_tokens + overlay_len
-                                    ].detach().to("cpu").contiguous()
+                                    seg = (
+                                        v[
+                                            num_computed_tokens : num_computed_tokens
+                                            + overlay_len
+                                        ]
+                                        .detach()
+                                        .to("cpu")
+                                        .contiguous()
+                                    )
                                 except Exception:
                                     seg = v.detach().to("cpu").contiguous()
                                 req_info[k] = seg
@@ -380,7 +386,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
             model_kwargs_extra = {}
             # Pass per-request additional information map for this step (no concat)
             if per_req_additional_information:
-                model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
+                model_kwargs_extra["additional_information_by_req_id"] = (
+                    per_req_additional_information
+                )
             # Always pass per-request runtime additional_information (persisted in request state)
             try:
                 per_req_runtime_info = []
@@ -433,12 +441,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
             # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
             # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
             try:
-                if (
-                    isinstance(multimodal_outputs, dict)
-                    and (
-                        "additional_information_update" in multimodal_outputs
-                        or "additional_information_update_by_req_id" in multimodal_outputs
-                    )
+                if isinstance(multimodal_outputs, dict) and (
+                    "additional_information_update" in multimodal_outputs
+                    or "additional_information_update_by_req_id" in multimodal_outputs
                 ):
                     # Option A: list[dict] in batch order
                     updates_list = multimodal_outputs.get(
@@ -446,7 +451,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
                     )
                     if isinstance(updates_list, list):
                         for idx, upd in enumerate(updates_list):
-                            if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
+                            if not isinstance(upd, dict) or idx >= len(
+                                self.input_batch.req_ids
+                            ):
                                 continue
                             req_id = self.input_batch.req_ids[idx]
                             self._merge_additional_information_update(req_id, upd)
@@ -462,7 +469,10 @@ class GPUARModelRunner(OmniGPUModelRunner):
                                 continue
                             self._merge_additional_information_update(req_id, upd)
             except Exception as e:
-                logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}")
+                logger.error(
+                    f"Error merging for requests:{self.input_batch.req_ids} additional \
+                        information update: {e}, with the multimodal_outputs as {multimodal_outputs}"
+                )
             if not self.broadcast_pp_output:
 
                 # Common case.
@@ -624,7 +634,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             invalid_req_indices=invalid_req_indices,
             async_output_copy_stream=self.async_output_copy_stream,
         )
-        
+
     def _merge_additional_information_update(self, req_id: str, upd: dict) -> None:
         req_state = self.requests.get(req_id)
         if req_state is None:
@@ -646,4 +656,4 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 merged[k] = new_list
             else:
                 merged[k] = v
-        setattr(req_state, "additional_information_cpu", merged)
+        setattr(req_state, "additional_information_cpu", merged)

vllm_omni/worker/gpu_diffusion_model_runner.py

@@ -501,13 +501,16 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             else:
                 cudagraph_runtime_mode = _cg_mode
 
-            with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_tokens_after_padding,
-                num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_runtime_mode,
-                batch_descriptor=batch_descriptor,
+            with (
+                self.maybe_randomize_inputs(input_ids),
+                set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_tokens_after_padding,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    cudagraph_runtime_mode=cudagraph_runtime_mode,
+                    batch_descriptor=batch_descriptor,
+                ),
             ):
                 outputs = self.model(
                     input_ids=input_ids,

vllm_omni/worker/gpu_model_runner.py

@@ -545,14 +545,17 @@ class OmniGPUModelRunner(GPUModelRunner):
                 if num_tokens_across_dp is not None:
                     num_tokens_across_dp[:] = num_tokens_after_padding
 
-            with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_tokens_after_padding,
-                num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_runtime_mode,
-                batch_descriptor=batch_descriptor,
-                ubatch_slices=ubatch_slices,
+            with (
+                self.maybe_randomize_inputs(input_ids),
+                set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_tokens_after_padding,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    cudagraph_runtime_mode=cudagraph_runtime_mode,
+                    batch_descriptor=batch_descriptor,
+                    ubatch_slices=ubatch_slices,
+                ),
             ):
                 outputs = self.model(
                     input_ids=input_ids,