[cpu] remove verbose print

Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>

.gitattributes

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

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

@@ -16,7 +16,7 @@ body:
       - Passwords or authentication credentials
       - Private URLs or endpoints
       - Personal or confidential data
-      
+
       Consider redacting or replacing sensitive values with placeholders like `<YOUR_TOKEN_HERE>` when sharing configuration or code examples.
 - type: textarea
   attributes:
@@ -58,7 +58,7 @@ body:
           input_modalities=["text"],
           output_modalities=["text"]
       )
-      
+
       dit_config = create_dit_stage_config(
           stage_id=1,
           model_path="stabilityai/stable-diffusion-2-1",

.gitignore

@@ -20,6 +20,7 @@ parts/
 sdist/
 var/
 wheels/
+vllm/
 share/python-wheels/
 *.egg-info/
 .installed.cfg
@@ -176,7 +177,6 @@ tests/
 
 # vLLM-omni specific
 # Model files and checkpoints
-models/
 checkpoints/
 *.bin
 *.safetensors

.pre-commit-config.yaml

@@ -0,0 +1,29 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      # list of supported hooks: https://pre-commit.com/hooks.html
+      - id: check-yaml
+        args: ["--unsafe"]
+      - id: debug-statements
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+        args: ["--fix=lf"]
+      - id: trailing-whitespace
+        args: ["--markdown-linebreak-ext=md"]
+
+  - repo: https://github.com/PyCQA/isort
+    rev: 5.13.2
+    hooks:
+      - id: isort
+
+  - repo: https://github.com/psf/black
+    rev: 24.10.0
+    hooks:
+      - id: black
+
+  - repo: https://github.com/PyCQA/flake8
+    rev: 7.1.1
+    hooks:
+      - id: flake8
+        args: ["--max-line-length=88", "--extend-ignore=E203,W503"]

CONTRIBUTING.md

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

README.md

@@ -38,26 +38,71 @@ vLLM-omni is built on a modular architecture that extends vLLM's core functional
 
 ## 🛠️ Installation
 
-Set up basic environments
+### 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
 ```
-Install certain version of vllm with commitid: 808a7b69df479b6b3a16181711cac7ca28a9b941
+
+### 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 808a7b69df479b6b3a16181711cac7ca28a9b941
-VLLM_USE_PRECOMPILED=1 uv pip install --editable .
+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 -r requirements.txt
 ```
 
 ## Run examples (Qwen2.5-omni)
 
 Get into the example folder
 ```bash
-cd vllm_omni
-cd examples/offline_inference/qwen2_5_omni
+cd examples/offline_inference/qwen_2_5_omni
 ```
 Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run.
 ```bash
@@ -65,11 +110,6 @@ bash run.sh
 ```
 The output audio is saved in ./output_audio
 
-## To-do list
-- [x] Offline inference example for Qwen2.5-omni with single request
-- [ ] Adaptation from current vllm branch to stable vllm v0.11.0
-- [ ] Offline inference example for Qwen2.5-omni with streaming multiple requests
-- [ ] Online inference support
-- [ ] Support for other models
+## Further details
 
-For detailed model management, see [vllm_omni_design.md](docs/architecture/vllm_omni_design.md) and [high_level_arch_design.md](docs/architecture/high_level_arch_design.md).
+For detailed architecture design, see [vllm_omni_design.md](docs/architecture/vllm_omni_design.md) and [high_level_arch_design.md](docs/architecture/high_level_arch_design.md).

docs/api/API_DESIGN_TEMPLATE.md

@@ -73,20 +73,20 @@ from enum import Enum
 
 class DiTCacheManager:
     """Manages DiT cache for diffusion models."""
-    
+
     def __init__(self, config: DiTCacheConfig):
         self.config = config
         self.cache = {}
         self.cache_stats = {}
-    
+
     async def get_cache(self, cache_key: str) -> Optional[Any]:
         """Retrieve cached data."""
         pass
-    
+
     async def set_cache(self, cache_key: str, data: Any) -> None:
         """Store data in cache."""
         pass
-    
+
     async def invalidate_cache(self, cache_key: str) -> None:
         """Invalidate cached data."""
         pass
@@ -103,7 +103,7 @@ class OmniScheduler():
         self.scheduler = self._create_scheduler()
         self.cache_manager = DiTCacheManager(config.dit_cache_config)
         self._running = False
-    
+
     def _create_scheduler(self) -> BaseScheduler:
         """Factory method to create appropriate scheduler."""
         if self.config.scheduler_type == SchedulerType.FIFO:
@@ -118,13 +118,13 @@ class OmniScheduler():
     async def schedule(self, request: OmniRequest) -> bool:
         """
         Schedule a request for processing.
-        
+
         Args:
             request: The request to schedule
-            
+
         Returns:
             bool: True if successfully scheduled, False otherwise
-            
+
         Raises:
             SchedulerError: If scheduling fails
             QueueFullError: If queue is at capacity
@@ -137,7 +137,7 @@ class OmniScheduler():
     def update_config(self, new_config: SchedulerConfig) -> None:
         """Update scheduler configuration."""
         pass
-    
+
     def get_config(self) -> SchedulerConfig:
         """Get current configuration."""
         pass
@@ -149,12 +149,12 @@ class OmniScheduler():
         """Start the scheduler."""
         self._running = True
         # Start background tasks
-    
+
     async def stop(self) -> None:
         """Stop the scheduler gracefully."""
         self._running = False
         # Cleanup resources
-    
+
     async def shutdown(self) -> None:
         """Force shutdown the scheduler."""
         # Immediate cleanup
@@ -170,7 +170,7 @@ class OmniScheduler():
             "processed_requests": self.scheduler.processed_count(),
             "cache_hit_rate": self.cache_manager.get_hit_rate()
         }
-    
+
     def get_metrics(self) -> Dict[str, float]:
         """Get performance metrics."""
         pass
@@ -182,21 +182,21 @@ class OmniScheduler():
 @dataclass
 class OmniConfig:
     """Configuration for the core module."""
-    
+
     # Scheduler configuration
     scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)
-    
+
     # Cache configuration
     dit_cache: DiTCacheConfig = field(default_factory=DiTCacheConfig)
-    
+
     # Resource limits
     max_memory_gb: float = 16.0
     max_gpu_utilization: float = 0.8
-    
+
     # Timeouts
     request_timeout: int = 300
     worker_timeout: int = 60
-    
+
     def validate(self) -> None:
         """Validate configuration parameters."""
         if self.max_memory_gb <= 0:

docs/api/README.md

@@ -157,10 +157,10 @@ Each API template follows this standardized structure:
    # Clone the repository
    git clone https://github.com/hsliuustc0106/vllm-omni.git
    cd vllm-omni
-   
+
    # Install dependencies
    pip install -r requirements-dev.txt
-   
+
    # Set up pre-commit hooks
    pre-commit install
    ```

docs/architecture/high_level_arch_design.md

@@ -52,7 +52,7 @@ According to analysis for current popular open-source models, most of them have
 * **Extension, multiple AR+DiT (e.g. Qwen-omni)**  
   A new model for multimodal input and output. It designs the Thinker-Talker-Codec structure, which is a double-AR+DiT format.  
   ![][image5]  
-  
+
 
 
 # Architectural Overview
@@ -112,8 +112,8 @@ vLLM-omni will import core modules from existing vLLM framework and introduce se
 
 # Contact
 
-For any questions or feedback regarding this design document, please contact   
-Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk),   
+For any questions or feedback regarding this design document, please contact  
+Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk),  
 Hongsheng Liu: [hsliuustc@gmail.com](mailto:hsliuustc@gmail.com)
 
 [image1]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAALYAAAFwCAYAAAAVNoEZAABxw0lEQVR4XuydBVRcWb6ve90709OStnTScel4x93dIUaEuBJ3dzciRPDgEtzdCRpCBCIEEkIgQHCCBHe+t6uSvNvNnZ55r2/fng5zvrX+69TZp4qkan3nz2/vKup88sm/OWqXL3HxwnnOnj7JubOnuaB6jqtXLqJ28TxXRWmrX8Ha0gIzE0OuX1bFQFsTcxMDLMyMMDXS44bWVRxsLLG5aYKuxlVsLMxwc3bAWF8bNzdnHGwt8fb2wNXZHlcnW3HMDjsbC4Ju3cLO2hwvT3ecnexwsLcmIMAHHy83+X18PVxxtLMUt+1xcbTBw8URRwdr+f1tLM1wd3PB08NZlAs24ue4uTrKf0bD5yfxb8qWDWvZtnE9W0Vt37yBvTu3cWDPTlE7OHviCBfPnsTEwECIfg7Vk0e5fkkVoxtamBneQF1NlSsXTmMrZJaN6elcF7LexFaIZ6SngbOjvfy2l5DUSwhoa2mK9U1jfL09ibwTKaS3IPx2BA4ONnh7uuHr4yUX1d/Pi0B/X1xd7PHz8cTT3UUc85SfIMHBQUJ4a3EfH/mJ4OsjThoXB/njPNydJbEl3rFisTKrli5i9fLFrF21jC0bVdi2aR07tqzn+KF9nD91HF1NDU4eO8CZIwe5fO4UBjoacpEvnDnG5fMnsDQ1Rk/rOje0rwmxLeQd3chAC0d7W6zMjfBwc8JHiGlmpIu5oS6hISHcibgtOrE99+4+wE50dR8hraeHm1zOW4G+BAb4yjtyYICf/FiAED0wMICw0DDxODvR8QNEBeLv7yO6txPeotN7e7lLYku8Y8akccycOhGlGVNRVprJsoXzWLlEmbUrl7Jv5xbOHD8sBD/ApnWrOLZvN2ePH0FDRJVronMf3rdDiH1aiK4p9s+JE0AWS6xE59bA2EAHO6ub3DTRw1UI7Ovjg/oVEWV01Llz567ouH6i+wbz8FEMtiKaBPr7yzuvn6+nED9I3rH9/bwJDgqUSy3r1OHh4YSEBOPu7io6/W1CQ0PFSeAvxHeTd3B/X19JbIl3TB4zjOnjRzJz8jiUFCazYPZ0li1QYrXo4nt3bObE4QPs3LIRlRWLObh7h9jfj5rqGdGtj7N/11Yh9hk0r14W+yfQ0bgiZLZC8/pljPR1ROwwwdxYF2d7GyGfN1cunkJfiB0h4oe3lweRkfeIjn4iF9vf1wcnkbX9fL0ICb4ll1kutrgtkzdEyB4ZGUlQ0C35Y+/cuSPv3oHimOxkkHX2W6KjN3x+Ev+mjBzQmzGD+zJxxCAmjxrC9AmjUBZyL1dWEkJvEF17G6tEB1++aB47Nq0Xsm+Vy71PbHdv2yTP4LL9Q3t3onVdTYhsyCXV02LyqCuyuS762texlk0oXV05e/IQetrqBInuK5v8RUc/JiLiDvYiisgmf7Y2N0UssZJv1a9fFQJ7yiUPkXXrsBAePHiAn+jMsghy//590dlDRRTxk8cW2VhoWLAktsQ7Bv7UleH9egi5+8kFl3Xw+TOnsmTBLDGpXMsOMaFcunCuiChKbF27ml1CbtnEcvt6FbnY504eE517uxB+CxpXLwlxNbhwXjbhvCGXWFOMmZsZibztwMmj+7khsniA6Kwurs4iSoi87OKCuuj4urqa3DQ3xdHBVt6pXZwd0NJQJ1AmbGiw6NC3xYkQLbK6lxA6WC55yP8VWxZrbhEWHiKJLfGOPl06MKhnVwb06Cy2XZg0aqiIJJNYNHcG61cvl2dtZXF78byZrF22iHUrl7B53WqRwxcI6dfLO/WW9Wvk3f3i2VOonj7BeRFT9HW0ULtwRoh/FIMb2hgZ6otMLuvql9klos3COYri3xrOiEH96dujK907d+Snrp3o37c3Gzdu4MKFCxjqG7Bn9y7uRt6Vd+gHUQ9xFsKHiMmnbEy29ReRRbaaEhYWJo8qDZ+fxL8pfbt2pF+3TqJ+ZHCvLiKODGX2tAly8VRWLJFPJBcoKcpFX6E8j9WL300slyrPkXd0WdfeIE4A2UqKbHlQtpJy7tQxtNWvckZILYspOprX0dFSZ8+2DezcqMKsCaOZOGwQvTp1oHPbVnRo2YL2rZrz/TdN+O6rJrRp8YMQvA+q5y9w9NARbgUFEf3wEffuR2FrZyUXOiIiQnTyEPmkUSa2LNLcu/dAElviHb06taF357bIBB/cqzNjh/YXOXs0cxWnsWjebPlKyVwRTebNmILyLAWWzJ3NMuW5LJijIO/o61YvY/WShWxUWcmB3dvZLgQ/emAvZ04cZc/2zULmTaiKuLJv22ZmjxvJqH496dK6pagW9Gnfhp7tWtH5h+9p//23tPruG1qI+v6br/n2qy/47K9/odn337FGRQUHJ2fCbt/GyFhPxI4Q+eRRJrivr7d8AvngQZTo6tGS2BLv6P6jELtLO7nY/bv/yPCBPZkiBJylMIV5QuS5M6czR3EKStMnMU8m+5yZLFGazdwZU1mzfAmrly2UTyzXr1omjyMb1qxgn8jch/fvYdsGFXZsXMMusV0i7j+6V1eGdm3PGBF5pvXtztxBPZjetyujxL8/VPw/erRuzo9C8pZNheBNv+b7r7/iy88/o2PHjixcspTbtyO4IbK4bMnvzp13KySyzB0cIptYRktiS/wXnzdpzpdNmtGtQyt+6tiK/j06MXbEYCaPG8WMKROYOWUc0yaOY6ro4orittKMacxWmMysaeNZtmAei+fPEZ19FquWLmSdkHv5ovnyzL1dxJQ1Qvr1ohaOG8rxJTO4ukYJ/a1LODprDEsGdmdkux8Y2KopHb/+knZffk6LJp/xvRD5e3G7+Zdf0PKbr2gtJG8jZP/uqy9Zu24tqmdOEn77jrxjy96wkb1BExERyd2794m4c1cSW+Idn33ZlL99IaJA6x/o1KaZ6N7tGTmkP+NGDpXLPG38KLnkk8eNYPqkMXKpZ0ydIKQfK2LJHHkHV1aaIYReIM/jMtFlgss691KRzeeKn6G7YwVup3fgr3YAw83KLO/VkbX9OrNhYBfmd23JrC6tmNT+B7p89zVdRafu1uwbOjf9ivbffUUHsS/r4h3F/69Fs2YsWbhAZOn78owdGBgol1vWyWWTyYg7EZLYEu9o/cVXrO71E6aL53Nx1mT6iLw9sE9Phg/qx/jhQ+Q1doTYjhryXu6xTJNtx4scrjCVOSKiKImoIsvispJ19MULlEQGn8GsYf05vWI2+aGOZHqb4LJvNY+v7iP2+hEeXdpL2NF1uG9V5ti4nuwb1pUlPVoz88cWTGzblHGihrVuypA23zOsfXOGCPE7t/qBZt9+g7qmFmFh4XKxZXlb9kbPbdG1ZZ87afj8JP5NGd2mFfuHD8BCeQ43F8/hzMxxzB3al5H9fmJo/z4MH9CXEYMHMEpMKscNFxFl9AjGjxjERNHRFSbIJB/N9MljmDV9MjOnTUJh8nhmT5/CxNHDMNu/luI77pRGB5Dhb0t1zB2qnoRREuxEnpsx989sI3T/apzXKXBz6WRMl01Bf9F4VKcP5uj4Pmwd1p21g7ui3Lsj83p1oEerZrQV3fv7pt9y4uRxbgUHERIWRlh4BMFC9FvBoZLYEu/YObgfp0YO4ZLowtemjhViKXB1wVT2KIxBccRAxorOPVgIPnRQX8YMHSjv4GPeSz5p7EgmjhkhjywyoadPkmXxMfJMPmHoAOJstKhLjKLiaTiVCQ+pSY6jJj6K0jA3CjxNSdI/S4L2MR5c2sW9CzsIOqKC156lGC2bivr8MRye0I+dI3uyelBXlvXvxMA2zWnb7Ds+//QvjBkziiAxaQwJDRWTx3D5kqC/9Ja6xAcU27Rkcgvx67/59yzo2JrDo/thvVYZj51rCT20leCDm7DfvALN5XNRHj2UcQP7MKRPV9HJe4suLuvkQvZRw+SCTxg9XH57pDgZji6bSXVcJHWZidS+SaMuN43alGdysasfBlER6UNpsCNv/azJctQh1eoacdrHuX9hF67bFmK+SpFz04dwZHw/to3sxfqh3ZgssniXZl/LlwT/+p//ia+vn5g8+uPt7SX/YJSrq4sktsQ7un39Fd2/+ZoB337NpJZNWdezI/pKY3FSmU+o2mVCjc3xd/LCw9kHbSNrDl/QQFFpAaNHDGNov3dZfIzo3rLJ5tiR7/L4zJGDCBQ5ui4tntr8TOpF1WWnUJsUQ+2z+1RHBVF1L4DSUFeKAuzJdjEkzUaT+BtniLq8F+ftypitVOCsEPvoxH7sHt2XzUN7MLtHGwa0aEpLIfYXf/sMS0sLPDw8cHR0xN7ORl4Nn5/Evymd2rXgxzY/0EF0wj5dO3NE2xTj5CLMC+rQL6zDtKQe+/I6rEvrMS6owfBNFQY5leimlnLI0pOl2/aLDN6X0UPede9TG1dxS/sMiO5cHO5DcYgHed4OZNkZURbqSfWDYKrvBlN+25/iIDcK/Rx542JEpp0uTzWPEXl+Jw6b52O8fDrnFYZxcupg9o7pK7p2Dxb3aYdil5a0+eYrvvnicw4e3I/lTQsM9G+gr6eL3g1dSWyJd3T44Rt+bN2C+bsOc9w7ApMiuFkMN/JrsSyuwbm8FuuiWkzF/s3Caizf1mCZX4VpbgV6yW+5/iCRpVv3MkFMJBdMn0BehDf50SGQLaLH6xfUy+JHVCiVYZ6UeltRGexO9T2xHxFI0S2RtX0cyHMxI9P2Bo9Flw8/vQ3bjfMwWDqVc0LsE1MGsXd0H7aP6M6SPu2Z3bUVbb9pwjdffsnKlcvQ1rjO1SuXuHZVVpclsf8gNomSvdgfKuaXh//1qD3O4fLrKjRya9HJrUE3uxqr4mrcy2vwLKvDMKsKSyG0s9i3Lxad+20txqJj3xRlJeQ2yyhBNy6Dq0GPWLN8Kbn3g6jNEdEjL5uHphZcO3QRNyNHEnyCqY2NpfbFC2rjnlBxL0x0clkMuUmSmQZPdS9wR3U/t4Line truncated

docs/architecture/implementation_architecture.md

@@ -1,353 +0,0 @@
-# vLLM-omni Implementation Architecture
-
-## 1. Package Structure
-
-```
-vllm_omni/
-├── __init__.py                 # Main package exports
-├── config/                     # Configuration management
-│   ├── __init__.py
-│   ├── stage_config.py        # OmniStageConfig implementation
-│   ├── dit_config.py          # DiT-specific configurations
-│   └── cache_config.py        # Caching configurations
-├── core/                       # Core processing components
-│   ├── __init__.py
-│   ├── stage_manager.py       # Multi-stage orchestration
-│   ├── dit_cache_manager.py   # DiT caching system
-│   └── sched/                 # Schedulers
-│       ├── __init__.py
-│       ├── scheduler.py       # Base scheduler interface
-
-├── engine/                     # Engine components
-│   ├── __init__.py
-│   ├── processor.py           # Output processing
-│   └── output_processor.py    # Multimodal output handling
-├── executor/                   # Executor implementations
-│   ├── __init__.py
-│   ├── base_executor.py       # Base executor interface
-│   └── diffusers_executor.py  # Diffusers pipeline executor
-├── model_executor/            # Model execution components
-│   ├── __init__.py
-├── worker/                     # Worker implementations
-│   ├── __init__.py
-│   ├── gpu_diffusion_model_runner.py  # Loads and runs diffusers pipelines
-│   └── gpu_diffusion_worker.py        # Worker facade around the diffusion runner
-├── entrypoints/               # Entry points and CLI
-│   ├── __init__.py
-│   ├── omni.py               # OmniServeCommand
-│   ├── omni_llm.py            # OmniLLM and AsyncOmniLLM
-│   └── cli/                  # CLI integration
-│       ├── __init__.py
-│       └── main.py           # CLI main entry point
-├── request.py                 # Request handling
-├── dit_cache_interface.py     # DiT cache interface
-└── utils/                     # Utility functions
-    ├── __init__.py
-    ├── multimodal.py         # Multimodal utilities
-    └── vae.py                # VAE utilities for image processing
-```
-
-## 2. Core Module Dependencies
-
-### 2.1 vLLM Integration Points
-
-```python
-# Key vLLM imports and extensions
-from vllm.v1.engine.llm_engine import LLMEngine
-from vllm.v1.engine.async_llm import AsyncLLM
-from vllm.v1.core.sched.scheduler import SchedulerInterface
-from vllm.v1.executor.multiproc_executor import MultiprocExecutor
-from vllm.v1.worker.gpu_worker import Worker as GPUWorker
-from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-from vllm.v1.outputs import ModelRunnerOutput
-```
-
-### 2.2 Internal Dependencies
-
-```python
-# Internal module dependencies
-vllm_omni.config → vllm_omni.core
-vllm_omni.core → vllm_omni.engine, vllm_omni.executor, vllm_omni.worker
-vllm_omni.engine → vllm_omni.model_executor
-vllm_omni.entrypoints → vllm_omni.core
-```
-
-## 3. Configuration System
-
-### 3.1 Stage Configuration
-
-```python
-# vllm_omni/config/stage_config.py
-@dataclass
-class OmniStageConfig:
-    stage_id: int
-    engine_type: Literal["AR", "DiT"]
-    model_path: str
-    input_modalities: List[str]
-    output_modalities: List[str]
-    vllm_config: Optional[VllmConfig] = None
-    dit_config: Optional[DiTConfig] = None
-    executor_class: Optional[type[Executor]] = None
-    stage_output: Optional[Any] = None
-```
-
-### 3.2 DiT Configuration
-
-```python
-# vllm_omni/config/stage_config.py
-@dataclass
-class DiTConfig:
-    model_config: Optional[ModelConfig] = None
-    scheduler_config: Optional[Any] = None
-    device_config: Optional[DeviceConfig] = None
-    load_config: Optional[LoadConfig] = None
-    compilation_config: Optional[CompilationConfig] = None
-    dit_cache_config: Optional[DiTCacheConfig] = None
-    num_inference_steps: int
-    guidance_scale: float = 7.5
-    use_diffusers: bool = False
-    diffusers_pipeline: Optional[str] = None
-    height: int = 512
-    width: int = 512
-```
-
-### 3.3 Cache Configuration
-
-```python
-# vllm_omni/config/cache_config.py
-@dataclass
-class DiTCacheConfig:
-    cache_tensors: List[DiTCacheTensor]
-    max_cache_size: int
-    cache_strategy: str = "fifo"
-    enable_optimization: bool = True
-```
-
-## 4. Core Implementation Details
-
-### 4.1 OmniLLM Implementation
-
-```python
-# vllm_omni/core/omni_llm.py
-class OmniLLM(LLM):
-    def __init__(self, stage_configs: List[OmniStageConfig]):
-        super().__init__()
-        self.stage_configs = stage_configs
-        self.engine_list: List[LLMEngine] = []
-        self.output_processor = MultimodalOutputProcessor()
-        self._initialize_stage_engines()
-    
-    def _initialize_stage_engines(self) -> None:
-        """Initialize LLMEngine instances for each stage"""
-        for stage_config in self.stage_configs:
-            if stage_config.engine_type == "AR":
-                engine = self._create_ar_engine(stage_config)
-            elif stage_config.engine_type == "DiT":
-                engine = self._create_dit_engine(stage_config)
-            self.engine_list.append(engine)
-    
-    def generate(self, stage_args_list: List[Dict], **kwargs) -> List[RequestOutput]:
-        """Main generation interface - orchestrates multi-stage processing"""
-        current_output = None
-        
-        for i, (stage_config, stage_args) in enumerate(zip(self.stage_configs, stage_args_list)):
-            stage_engine = self.engine_list[i]
-            
-            # Prepare input for this stage
-            processed_input = self._process_stage_inputs(stage_config, stage_args, current_output)
-            
-            # Execute stage
-            stage_output = self._execute_stage(stage_engine, processed_input)
-            
-            # Update for next stage
-            current_output = stage_output
-            stage_config.stage_output = stage_output
-        
-        # Process final output
-        return self.output_processor.process_output(current_output)
-```
-
-### 4.2 AsyncOmniLLM Implementation
-
-```python
-# vllm_omni/core/omni_llm.py
-class AsyncOmniLLM(AsyncLLM):
-    def __init__(self, stage_configs: List[OmniStageConfig]):
-        super().__init__()
-        self.stage_configs = stage_configs
-        self.async_engine_list: List[AsyncLLM] = []
-        self.output_processor = MultimodalOutputProcessor()
-        self._initialize_async_stage_engines()
-    
-    async def generate_async(self, stage_args_list: List[Dict], **kwargs) -> List[RequestOutput]:
-        """Async generation interface"""
-        # Similar to OmniLLM but with async/await patterns
-        pass
-``
-
-### 4.3 Model Runners
-
-```python
-# vllm_omni/model_executor/dit_model_runner.py
-class OmniDiffusionModelRunner(GPUModelRunner):
-    def execute_model(
-        self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-    ) -> Union[ModelRunnerOutput, IntermediateTensors]:
-        # DiT model execution logic
-        return ModelRunnerOutput(
-            req_ids=[...],
-            req_id_to_index={...},
-            sampled_token_ids=[],
-            spec_token_ids=None,
-            logprobs=None,
-            prompt_logprobs_dict={},
-            pooler_output=[tensor, ...],  # DiT output tensors
-            kv_connector_output=None,
-            num_nans_in_logits=None,
-        )
-
-# vllm_omni/model_executor/ar_model_runner.py
-class OmniARModelRunner(GPUModelRunner):
-    def execute_model(
-        self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-    ) -> Union[ModelRunnerOutput, IntermediateTensors]:
-        # AR model execution with hidden state output
-        return ModelRunnerOutput(
-            req_ids=[...],
-            req_id_to_index={...},
-            sampled_token_ids=[],
-            spec_token_ids=None,
-            logprobs=None,
-            prompt_logprobs_dict={},
-            pooler_output=[hidden_states, ...],  # Hidden states
-            kv_connector_output=None,
-            num_nans_in_logits=None,
-        )
-```
-
-### 4.5 Output Processing
-
-```python
-# vllm_omni/engine/output_processor.py
-class MultimodalOutputProcessor(OutputProcessor):
-    def __init__(self):
-        self.output_handlers: Dict[str, Callable] = {
-            "image": self._process_image_output,
-            "text+image": self._process_text_image_output,
-            "latents": self._process_latents_output,
-            "text": self._process_text_output,
-        }
-    
-    def process_outputs(self, engine_core_outputs: List[EngineCoreOutput], ...):
-        """Process multimodal outputs based on type"""
-        for engine_core_output in engine_core_outputs:
-            output_type = self._detect_output_type(engine_core_output)
-            handler = self.output_handlers.get(output_type, self._process_pooling_output)
-            handler(engine_core_output)
-```
-
-## 5. CLI Integration
-
-### 5.1 Entry Point Override
-
-```python
-# vllm_omni/entrypoints/cli/main.py
-def main():
-    """Main CLI entry point that intercepts vLLM commands"""
-    if "--omni" in sys.argv:
-        omni_args = [arg for arg in sys.argv[1:] if arg != "--omni"]
-        omni_serve = OmniServeCommand()
-        omni_serve.run(omni_args)
-    else:
-        from vllm.entrypoints.cli.main import main as vllm_main
-        vllm_main()
-```
-
-### 5.2 Package Configuration
-
-```toml
-# pyproject.toml updates
-[project.scripts]
-vllm = "vllm_omni.entrypoints.cli.main:main"
-vllm-omni = "vllm_omni.entrypoints.cli.main:main"
-
-[project.entry-points."vllm.plugins"]
-omni = "vllm_omni.plugin:OmniPlugin"
-```
-
-## 6. Testing Strategy
-
-### 6.1 Unit Tests
-- Individual component testing
-- Mock vLLM dependencies
-- Configuration validation
-
-### 6.2 Integration Tests
-- End-to-end pipeline testing
-- vLLM compatibility testing
-- Multi-stage processing validation
-
-### 6.3 Performance Tests
-- Benchmarking against native vLLM
-- Memory usage profiling
-- Latency measurements
-
-## 7. Installation and Setup
-
-### 7.1 Package Installation
-```bash
-pip install vllm>=0.10.2
-pip install vllm-omni
-```
-
-### 7.2 Development Setup
-```bash
-git clone https://github.com/hsliuustc0106/vllm-omni
-cd vllm-omni
-pip install -e ".[dev]"
-```
-
-### 7.3 Usage
-```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8000
-```
-
-## 8. Memory Management
-
-### 8.1 DiT Cache Management
-- Tensor caching for intermediate results
-- Memory pooling for efficient allocation
-- Cache eviction strategies
-
-### 8.2 Multi-Stage Memory
-- Inter-stage data passing optimization
-- Memory sharing between stages
-- Garbage collection optimization
-
-## 9. Error Handling
-
-### 9.1 Stage Failure Handling
-- Graceful degradation on stage failures
-- Error propagation and reporting
-- Recovery mechanisms
-
-### 9.2 vLLM Compatibility
-- Version compatibility checks
-- API change detection
-- Fallback mechanisms
-
-## 10. Monitoring and Logging
-
-### 10.1 Performance Metrics
-- Stage execution times
-- Memory usage per stage
-- Throughput measurements
-
-### 10.2 Debug Information
-- Request tracing across stages
-- Cache hit/miss ratios
-- Error logging and reporting

docs/architecture/vllm_omni_design.md

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

docs/mrs_design.md

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

examples/basic/api_client.py

@@ -1,128 +0,0 @@
-#!/usr/bin/env python3
-"""
-Simple API client example for vLLM-omni server.
-
-This example shows how to interact with the vLLM-omni API server
-using HTTP requests.
-"""
-
-import requests
-import json
-import time
-
-
-def test_health(host="localhost", port=8000):
-    """Test if the server is healthy."""
-    url = f"http://{host}:{port}/health"
-    try:
-        response = requests.get(url, timeout=5)
-        if response.status_code == 200:
-            print("✅ Server is healthy")
-            return True
-        else:
-            print(f"❌ Server health check failed: {response.status_code}")
-            return False
-    except requests.exceptions.RequestException as e:
-        print(f"❌ Cannot connect to server: {e}")
-        return False
-
-
-def generate_text(prompt, host="localhost", port=8000, **kwargs):
-    """Generate text using the API server."""
-    url = f"http://{host}:{port}/generate"
-    
-    # Default parameters
-    params = {
-        "prompts": [prompt],
-        "max_tokens": kwargs.get("max_tokens", 100),
-        "temperature": kwargs.get("temperature", 0.7),
-        "top_p": kwargs.get("top_p", 1.0),
-        "frequency_penalty": kwargs.get("frequency_penalty", 0.0),
-        "presence_penalty": kwargs.get("presence_penalty", 0.0),
-    }
-    
-    # Add optional parameters
-    if "stop" in kwargs:
-        params["stop"] = kwargs["stop"]
-    
-    try:
-        response = requests.post(url, json=params, timeout=30)
-        response.raise_for_status()
-        return response.json()
-    except requests.exceptions.RequestException as e:
-        print(f"❌ API request failed: {e}")
-        return None
-
-
-def main():
-    """Main example function."""
-    print("vLLM-omni API Client Example")
-    print("=" * 40)
-    
-    # Test server health
-    if not test_health():
-        print("Please start the server first:")
-        print("  vllm serve Qwen/Qwen3-0.6B --omni --port 8000")
-        return
-    
-    # Example 1: Simple generation
-    print("\n1. Simple text generation:")
-    result = generate_text("Hello, how are you?", max_tokens=50)
-    if result:
-        text = result["outputs"][0]["text"]
-        print(f"   Input: Hello, how are you?")
-        print(f"   Output: {text}")
-    
-    # Example 2: Creative writing
-    print("\n2. Creative writing:")
-    result = generate_text(
-        "Write a short story about a robot learning to paint",
-        max_tokens=150,
-        temperature=0.9
-    )
-    if result:
-        text = result["outputs"][0]["text"]
-        print(f"   Input: Write a short story about a robot learning to paint")
-        print(f"   Output: {text}")
-    
-    # Example 3: Question answering
-    print("\n3. Question answering:")
-    result = generate_text(
-        "What is the capital of France?",
-        max_tokens=30,
-        temperature=0.3
-    )
-    if result:
-        text = result["outputs"][0]["text"]
-        print(f"   Input: What is the capital of France?")
-        print(f"   Output: {text}")
-    
-    # Example 4: Multiple prompts
-    print("\n4. Multiple prompts:")
-    prompts = [
-        "Tell me a joke",
-        "What's 2+2?",
-        "Describe the color blue"
-    ]
-    
-    for prompt in prompts:
-        result = generate_text(prompt, max_tokens=30)
-        if result:
-            text = result["outputs"][0]["text"]
-            print(f"   Q: {prompt}")
-            print(f"   A: {text}")
-    
-    # Example 5: Show full response structure
-    print("\n5. Full response structure:")
-    result = generate_text("Hello world", max_tokens=20)
-    if result:
-        print("   Full API response:")
-        print(json.dumps(result, indent=2))
-    
-    print("\n" + "=" * 40)
-    print("API client examples completed!")
-
-
-if __name__ == "__main__":
-    main()
-

examples/basic/text_generation.py

@@ -1,137 +0,0 @@
-#!/usr/bin/env python3
-"""
-Simple vLLM-omni usage example.
-
-This example demonstrates how to use vLLM-omni for basic text generation
-with a single AR (Autoregressive) stage.
-"""
-
-import asyncio
-from vllm_omni.entrypoints.omni_llm import OmniLLM, AsyncOmniLLM
-from vllm_omni.config import create_ar_stage_config
-
-
-def sync_example():
-    """Synchronous usage example."""
-    print("=== Synchronous vLLM-omni Example ===")
-    
-    # Create a simple AR stage configuration
-    stage_config = create_ar_stage_config(
-        stage_id=0,
-        model_path="Qwen/Qwen3-0.6B",
-        input_modalities=["text"],
-        output_modalities=["text"]
-    )
-    
-    # Initialize OmniLLM
-    omni_llm = OmniLLM([stage_config])
-    
-    # Prepare stage arguments
-    stage_args = [{
-        "prompt": "Hello, how are you today?",
-        "max_tokens": 50,
-        "temperature": 0.7
-    }]
-    
-    # Generate text
-    print(f"Input: {stage_args[0]['prompt']}")
-    outputs = omni_llm.generate(stage_args)
-    
-    # Display results
-    for i, output in enumerate(outputs):
-        print(f"Output {i}:")
-        if hasattr(output, 'outputs') and output.outputs:
-            for completion in output.outputs:
-                print(f"  Text: {completion.text}")
-                print(f"  Finished: {completion.finish_reason != 'length'}")
-                print(f"  Tokens: {len(completion.token_ids)} tokens")
-
-
-async def async_example():
-    """Asynchronous usage example."""
-    print("\n=== Asynchronous vLLM-omni Example ===")
-    
-    # Create a simple AR stage configuration
-    stage_config = create_ar_stage_config(
-        stage_id=0,
-        model_path="Qwen/Qwen3-0.6B",
-        input_modalities=["text"],
-        output_modalities=["text"]
-    )
-    
-    # Initialize AsyncOmniLLM
-    omni_llm = AsyncOmniLLM([stage_config])
-    
-    # Prepare stage arguments
-    stage_args = [{
-        "prompt": "What is artificial intelligence?",
-        "max_tokens": 100,
-        "temperature": 0.8
-    }]
-    
-    # Generate text asynchronously
-    print(f"Input: {stage_args[0]['prompt']}")
-    outputs = await omni_llm.generate_async(stage_args)
-    
-    # Display results
-    for i, output in enumerate(outputs):
-        print(f"Output {i}:")
-        if hasattr(output, 'outputs') and output.outputs:
-            for completion in output.outputs:
-                print(f"  Text: {completion.text}")
-                print(f"  Finished: {completion.finish_reason != 'length'}")
-                print(f"  Tokens: {len(completion.token_ids)} tokens")
-
-
-def multi_prompt_example():
-    """Example with multiple prompts."""
-    print("\n=== Multi-Prompt Example ===")
-    
-    # Create stage configuration
-    stage_config = create_ar_stage_config(
-        stage_id=0,
-        model_path="Qwen/Qwen3-0.6B",
-        input_modalities=["text"],
-        output_modalities=["text"]
-    )
-    
-    # Initialize OmniLLM
-    omni_llm = OmniLLM([stage_config])
-    
-    # Multiple prompts
-    prompts = [
-        "Tell me a joke",
-        "What's the weather like?",
-        "Explain quantum computing"
-    ]
-    
-    for prompt in prompts:
-        print(f"\nInput: {prompt}")
-        stage_args = [{
-            "prompt": prompt,
-            "max_tokens": 30,
-            "temperature": 0.9
-        }]
-        
-        outputs = omni_llm.generate(stage_args)
-        
-        if outputs and hasattr(outputs[0], 'outputs') and outputs[0].outputs:
-            completion = outputs[0].outputs[0]
-            print(f"Output: {completion.text}")
-
-
-if __name__ == "__main__":
-    print("vLLM-omni Simple Usage Examples")
-    print("=" * 40)
-    
-    # Run synchronous example
-    sync_example()
-    
-    # Run asynchronous example
-    asyncio.run(async_example())
-    
-    # Run multi-prompt example
-    multi_prompt_example()
-    
-    print("\n" + "=" * 40)
-    print("Examples completed!")

examples/offline_inference/qwen_2_5_omni/README.md

@@ -1,37 +1,92 @@
 # Offline Example of vLLM-omni for Qwen2.5-omni
 
-## Installation
+## 🛠️ Installation
 
-Set up basic environments
+### 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
 ```
-Install certain version of vllm with commitid: 808a7b69df479b6b3a16181711cac7ca28a9b941
+
+### 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 808a7b69df479b6b3a16181711cac7ca28a9b941
-VLLM_USE_PRECOMPILED=1 uv pip install --editable .
+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 .
 ```
 
-## Run examples
+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 -r requirements.txt
+```
+
+## 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 vllm_omni
-cd examples/offline_inference/qwen2_5_omni
+cd examples/offline_inference/qwen_2_5_omni
 ```
-Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run.
+Modify PYTHONPATH in run_multiple_prompts.sh as your path of vllm_omni. Then run.
 ```bash
-bash run.sh
+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
 ```
-The output audio is saved in ./output_audio
-
-## To-do list
-- [x] Offline inference example for Qwen2.5-omni with single request
-- [ ] Adaptation from current vllm branch to stable vllm v0.11.0
-- [ ] Offline inference example for Qwen2.5-omni with streaming multiple requests
-- [ ] Online inference support
-- [ ] Support for other models

examples/offline_inference/qwen_2_5_omni/end2end.py

@@ -1,18 +1,17 @@
 import argparse
 import os
-import soundfile as sf
+import os as _os_env_toggle
 import random
-import numpy as np
-import torch
 
+import numpy as np
+import soundfile as sf
+import torch
+from utils import make_omni_prompt
 from vllm.sampling_params import SamplingParams
 
-import os as _os_env_toggle
-_os_env_toggle.environ["VLLM_USE_V1"] = "1"
-
 from vllm_omni.entrypoints.omni_llm import OmniLLM
-from utils import make_omni_prompt
 
+_os_env_toggle.environ["VLLM_USE_V1"] = "1"
 
 SEED = 42
 # Set all random seeds
@@ -33,39 +32,123 @@ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--model', required=True, help='Path to merged model directory (will be created if downloading).')
-    parser.add_argument('--thinker-model', type=str, default=None)
-    parser.add_argument('--talker-model', type=str, default=None)
-    parser.add_argument('--code2wav-model', type=str, default=None)
-    parser.add_argument('--hf-hub-id', default='Qwen/Qwen2.5-Omni-7B', help='Hugging Face repo id to download if needed.')
-    parser.add_argument('--hf-revision', default=None, help='Optional HF revision (branch/tag/commit).')
-    parser.add_argument('--prompts', required=True, nargs='+', help='Input text prompts.')
-    parser.add_argument('--voice-type', default='default', help='Voice type, e.g., m02, f030, default.')
-    parser.add_argument('--code2wav-dir', default=None, help='Path to code2wav folder (contains spk_dict.pt).')
-    parser.add_argument('--dit-ckpt', default=None, help='Path to DiT checkpoint file (e.g., dit.pt).')
-    parser.add_argument('--bigvgan-ckpt', default=None, help='Path to BigVGAN checkpoint file.')
-    parser.add_argument('--dtype', default='bfloat16', choices=['float16', 'bfloat16', 'float32'])
-    parser.add_argument('--max-model-len', type=int, default=32768)
+    parser.add_argument(
+        "--model",
+        required=True,
+        help="Path to merged model directory (will be created if downloading).",
+    )
+    parser.add_argument("--thinker-model", type=str, default=None)
+    parser.add_argument("--talker-model", type=str, default=None)
+    parser.add_argument("--code2wav-model", type=str, default=None)
+    parser.add_argument(
+        "--hf-hub-id",
+        default="Qwen/Qwen2.5-Omni-7B",
+        help="Hugging Face repo id to download if needed.",
+    )
+    parser.add_argument(
+        "--hf-revision", default=None, help="Optional HF revision (branch/tag/commit)."
+    )
+    parser.add_argument(
+        "--prompts", nargs="+", default=None, help="Input text prompts."
+    )
+    parser.add_argument(
+        "--voice-type", default="default", help="Voice type, e.g., m02, f030, default."
+    )
+    parser.add_argument(
+        "--code2wav-dir",
+        default=None,
+        help="Path to code2wav folder (contains spk_dict.pt).",
+    )
+    parser.add_argument(
+        "--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt)."
+    )
+    parser.add_argument(
+        "--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file."
+    )
+    parser.add_argument(
+        "--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"]
+    )
+    parser.add_argument("--max-model-len", type=int, default=32768)
+    parser.add_argument(
+        "--init-sleep-seconds",
+        type=int,
+        default=20,
+        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
+    )
 
     parser.add_argument("--thinker-only", action="store_true")
     parser.add_argument("--text-only", action="store_true")
     parser.add_argument("--do-wave", action="store_true")
-    parser.add_argument('--prompt_type',
-                        choices=[
-                            'text', 'audio', 'audio-long', 'audio-long-chunks',
-                            'audio-long-expand-chunks', 'image', 'video',
-                            'video-frames', 'audio-in-video', 'audio-in-video-v2',
-                            "audio-multi-round", "badcase-vl", "badcase-text",
-                            "badcase-image-early-stop", "badcase-two-audios",
-                            "badcase-two-videos", "badcase-multi-round",
-                            "badcase-voice-type", "badcase-voice-type-v2",
-                            "badcase-audio-tower-1", "badcase-audio-only"
-                        ],
-                        default='text')
-    parser.add_argument('--use-torchvision', action='store_true')
-    parser.add_argument('--tokenize', action='store_true')
-    parser.add_argument('--output-wav', default="output.wav", help='Output wav file path.')
-    parser.add_argument('--thinker-hidden-states-dir', default="thinker_hidden_states", help='Path to thinker hidden states directory.')
+    parser.add_argument(
+        "--prompt_type",
+        choices=[
+            "text",
+            "audio",
+            "audio-long",
+            "audio-long-chunks",
+            "audio-long-expand-chunks",
+            "image",
+            "video",
+            "video-frames",
+            "audio-in-video",
+            "audio-in-video-v2",
+            "audio-multi-round",
+            "badcase-vl",
+            "badcase-text",
+            "badcase-image-early-stop",
+            "badcase-two-audios",
+            "badcase-two-videos",
+            "badcase-multi-round",
+            "badcase-voice-type",
+            "badcase-voice-type-v2",
+            "badcase-audio-tower-1",
+            "badcase-audio-only",
+        ],
+        default="text",
+    )
+    parser.add_argument("--use-torchvision", action="store_true")
+    parser.add_argument("--tokenize", action="store_true")
+    parser.add_argument(
+        "--output-wav", default="output.wav", help="[Deprecated] Output wav directory (use --output-dir)."
+    )
+    parser.add_argument(
+        "--output-dir", default="outputs", help="Output directory to save text and wav files together."
+    )
+    parser.add_argument(
+        "--thinker-hidden-states-dir",
+        default="thinker_hidden_states",
+        help="Path to thinker hidden states directory.",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
     args = parser.parse_args()
     return args
 
@@ -73,58 +156,98 @@ def parse_args():
 def main():
     args = parse_args()
     model_name = args.model
-    omni_llm = OmniLLM(model=model_name)
+    try:
+        # Preferred: load from txt file (one prompt per line)
+        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
+            with open(args.txt_prompts, "r", encoding="utf-8") as f:
+                lines = [ln.strip() for ln in f.readlines()]
+            args.prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
+    except Exception as e:
+        print(f"[Error] Failed to load prompts: {e}")
+        raise
+
+    if args.prompts is None:
+        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
+    omni_llm = OmniLLM(
+        model=model_name,
+        log_stats=args.enable_stats,
+        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+    )
     thinker_sampling_params = SamplingParams(
-                                            temperature=0.0,    # Deterministic - no randomness
-                                            top_p=1.0,          # Disable nucleus sampling
-                                            top_k=-1,           # Disable top-k sampling
-                                            max_tokens=2048,
-                                            seed=SEED,          # Fixed seed for sampling
-                                            detokenize=True,
-                                            repetition_penalty=1.1,
-                                            )
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
     talker_sampling_params = SamplingParams(
-                                            temperature=0.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,
-                                            stop_token_ids=[8294]
-                                            )
+        temperature=0.9,
+        top_p=0.8,
+        top_k=40,
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.05,
+        stop_token_ids=[8294],
+    )
     code2wav_sampling_params = SamplingParams(
-                                            temperature=0.0,    # Deterministic - no randomness
-                                            top_p=1.0,          # Disable nucleus sampling
-                                            top_k=-1,           # Disable top-k sampling
-                                            max_tokens=2048,
-                                            seed=SEED,          # Fixed seed for sampling
-                                            detokenize=True,
-                                            repetition_penalty=1.1,
-                                            )
+        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,
+    ]
 
-    sampling_params_list = [thinker_sampling_params,
-                            talker_sampling_params,
-                            code2wav_sampling_params]
-    
     prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
     omni_outputs = omni_llm.generate(prompt, sampling_params_list)
 
-    os.makedirs(args.output_wav, exist_ok=True)
+    # Determine output directory: prefer --output-dir; fallback to --output-wav
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    os.makedirs(output_dir, exist_ok=True)
     for stage_outputs in omni_outputs:
         if stage_outputs.final_output_type == "text":
             for output in stage_outputs.request_output:
-                request_id = output.request_id
+                request_id = int(output.request_id)
                 text_output = output.outputs[0].text
-                print(f"Request ID: {request_id}, Text Output: {text_output}")
+                # Save aligned text file per request
+                prompt_text = args.prompts[request_id]
+                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
             for output in stage_outputs.request_output:
-                request_id = output.request_id
+                request_id = int(output.request_id)
                 audio_tensor = output.multimodal_output["audio"]
-                output_wav = os.path.join(args.output_wav, f"output_{output.request_id}.wav")
-                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
+                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
+                sf.write(
+                    output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000
+                )
                 print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
 
 if __name__ == "__main__":
-    main()
+    main()

examples/offline_inference/qwen_2_5_omni/extract_prompts.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+import argparse
+from typing import Optional
+import os
+import torch
+
+
+def extract_prompt(line: str) -> Optional[str]:
+    # 提取第一个 '|' 与第二个 '|' 之间的内容
+    i = line.find('|')
+    if i == -1:
+        return None
+    j = line.find('|', i + 1)
+    if j == -1:
+        return None
+    return line[i + 1:j].strip()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", "-i", required=True, help="Input .lst file path")
+    parser.add_argument("--output", "-o", required=True, help="Output file path")
+    parser.add_argument("--topk", "-k", type=int, default=100, help="Extract the top K prompts (default: 100)")
+    args = parser.parse_args()
+
+    prompts = []
+    with open(args.input, "r", encoding="utf-8", errors="ignore") as f:
+        for line in f:
+            if len(prompts) >= args.topk:
+                break
+            p = extract_prompt(line.rstrip("\n"))
+            if p:
+                prompts.append(p)
+
+    with open(args.output, "w", encoding="utf-8") as f:
+        for p in prompts:
+            f.write(p + "\n")
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/offline_inference/qwen_2_5_omni/processing_omni.py

@@ -8,6 +8,7 @@ import time
 import warnings
 from functools import lru_cache
 from io import BytesIO
+
 import requests
 import torch
 import torchvision
@@ -16,7 +17,6 @@ from PIL import Image
 from torchvision import io, transforms
 from torchvision.transforms import InterpolationMode
 
-
 logger = logging.getLogger(__name__)
 
 IMAGE_FACTOR = 28
@@ -43,20 +43,24 @@ def round_by_factor(number: int, factor: int) -> int:
 
 
 def ceil_by_factor(number: int, factor: int) -> int:
-    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    """Returns the smallest integer greater than or equal to 'number' that is
+    divisible by 'factor'."""
     return math.ceil(number / factor) * factor
 
 
 def floor_by_factor(number: int, factor: int) -> int:
-    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    """Returns the largest integer less than or equal to 'number' that is
+    divisible by 'factor'."""
     return math.floor(number / factor) * factor
 
 
-def smart_resize(height: int,
-                 width: int,
-                 factor: int = IMAGE_FACTOR,
-                 min_pixels: int = MIN_PIXELS,
-                 max_pixels: int = MAX_PIXELS) -> tuple[int, int]:
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = IMAGE_FACTOR,
+    min_pixels: int = MIN_PIXELS,
+    max_pixels: int = MAX_PIXELS,
+) -> tuple[int, int]:
     """
     Rescales the image so that the following conditions are met:
 
@@ -68,7 +72,8 @@ def smart_resize(height: int,
     """
     if max(height, width) / min(height, width) > MAX_RATIO:
         raise ValueError(
-            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, "
+            f"got {max(height, width) / min(height, width)}"
         )
     h_bar = max(factor, round_by_factor(height, factor))
     w_bar = max(factor, round_by_factor(width, factor))
@@ -83,8 +88,9 @@ def smart_resize(height: int,
     return h_bar, w_bar
 
 
-def fetch_image(ele: dict[str, str | Image.Image],
-                size_factor: int = IMAGE_FACTOR) -> Image.Image:
+def fetch_image(
+    ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR
+) -> Image.Image:
     if "image" in ele:
         image = ele["image"]
     else:
@@ -105,10 +111,11 @@ def fetch_image(ele: dict[str, str | Image.Image],
         image_obj = Image.open(image)
     if image_obj is None:
         raise ValueError(
-            f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}"
+            f"Unrecognized image input, support local path, http url, "
+            f"base64 and PIL.Image, got {image}"
         )
     image = image_obj.convert("RGB")
-    ## resize
+    # resize
     if "resized_height" in ele and "resized_width" in ele:
         resized_height, resized_width = smart_resize(
             ele["resized_height"],
@@ -143,8 +150,10 @@ def smart_nframes(
             support either `fps` or `nframes`:
                 - nframes: the number of frames to extract for model inputs.
                 - fps: the fps to extract frames for model inputs.
-                    - min_frames: the minimum number of frames of the video, only used when fps is provided.
-                    - max_frames: the maximum number of frames of the video, only used when fps is provided.
+                    - min_frames: the minimum number of frames of the video,
+                      only used when fps is provided.
+                    - max_frames: the maximum number of frames of the video,
+                      only used when fps is provided.
         total_frames (int): the original total number of frames of the video.
         video_fps (int | float): the original fps of the video.
 
@@ -154,34 +163,39 @@ def smart_nframes(
     Returns:
         int: the number of frames for video used for model inputs.
     """
-    assert not ("fps" in ele
-                and "nframes" in ele), "Only accept either `fps` or `nframes`"
+    assert not (
+        "fps" in ele and "nframes" in ele
+    ), "Only accept either `fps` or `nframes`"
     if "nframes" in ele:
         nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
     else:
         fps = ele.get("fps", FPS)
-        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES),
-                                    FRAME_FACTOR)
+        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
         max_frames = floor_by_factor(
             ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
-            FRAME_FACTOR)
+            FRAME_FACTOR,
+        )
         nframes = total_frames / video_fps * fps
         nframes = min(max(nframes, min_frames), max_frames)
         nframes = round_by_factor(nframes, FRAME_FACTOR)
     if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
         raise ValueError(
-            f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}."
+            f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], "
+            f"but got {nframes}."
         )
     return nframes
 
 
-def _read_video_torchvision(ele: dict, ) -> torch.Tensor:
+def _read_video_torchvision(
+    ele: dict,
+) -> torch.Tensor:
     """read video using torchvision.io.read_video
 
     Args:
         ele (dict): a dict contains the configuration of video.
         support keys:
-            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
             - video_start: the start time of video.
             - video_end: the end time of video.
     Returns:
@@ -191,7 +205,8 @@ def _read_video_torchvision(ele: dict, ) -> torch.Tensor:
     if version.parse(torchvision.__version__) < version.parse("0.19.0"):
         if "http://" in video_path or "https://" in video_path:
             warnings.warn(
-                "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0."
+                "torchvision < 0.19.0 does not support http/https video path, "
+                "please upgrade to 0.19.0."
             )
         if "file://" in video_path:
             video_path = video_path[7:]
@@ -206,11 +221,10 @@ def _read_video_torchvision(ele: dict, ) -> torch.Tensor:
     total_frames, video_fps = video.size(0), info["video_fps"]
     total_duration = round(total_frames / video_fps, 3)
     logger.info(
-        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s"
+        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, "
+        f"duration={total_duration}s, time={time.time() - st:.3f}s"
     )
-    nframes = smart_nframes(ele,
-                            total_frames=total_frames,
-                            video_fps=video_fps)
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
     idx = torch.linspace(0, total_frames - 1, nframes).round().long()
     video = video[idx]
     return video, total_duration, nframes
@@ -222,34 +236,38 @@ def is_decord_available() -> bool:
     return importlib.util.find_spec("decord") is not None
 
 
-def _read_video_decord(ele: dict, ) -> torch.Tensor:
+def _read_video_decord(
+    ele: dict,
+) -> torch.Tensor:
     """read video using decord.VideoReader
 
     Args:
         ele (dict): a dict contains the configuration of video.
         support keys:
-            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
             - video_start: the start time of video.
             - video_end: the end time of video.
     Returns:
         torch.Tensor: the video tensor with shape (T, C, H, W).
     """
     import decord
+
     video_path = ele["video"]
     st = time.time()
     vr = decord.VideoReader(video_path)
     # TODO: support start_pts and end_pts
-    if 'video_start' in ele or 'video_end' in ele:
+    if "video_start" in ele or "video_end" in ele:
         raise NotImplementedError(
-            "not support start_pts and end_pts in decord for now.")
+            "not support start_pts and end_pts in decord for now."
+        )
     total_frames, video_fps = len(vr), vr.get_avg_fps()
     total_duration = round(total_frames / video_fps, 3)
     logger.info(
-        f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s"
+        f"decord:  {video_path=}, {total_frames=}, {video_fps=}, "
+        f"time={time.time() - st:.3f}s"
     )
-    nframes = smart_nframes(ele,
-                            total_frames=total_frames,
-                            video_fps=video_fps)
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
     idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
     video = vr.get_batch(idx).asnumpy()
     video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
@@ -272,27 +290,27 @@ def get_video_reader_backend() -> str:
         video_reader_backend = "decord"
     else:
         video_reader_backend = "torchvision"
-    # print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr)
+    # print(f"qwen-vl-utils using {video_reader_backend} to read video.",
+    #       file=sys.stderr)
     return video_reader_backend
 
 
 def fetch_video(
-        ele: dict,
-        image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
+    ele: dict, image_factor: int = IMAGE_FACTOR
+) -> torch.Tensor | list[Image.Image]:
     if isinstance(ele["video"], str):
         video_reader_backend = get_video_reader_backend()
-        video, total_dur, nframes = VIDEO_READER_BACKENDS[
-            video_reader_backend](ele)
+        video, total_dur, nframes = VIDEO_READER_BACKENDS[video_reader_backend](ele)
         frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes
         grid_timestamps = frame_timestamps[::FRAME_FACTOR]
         second_per_grid = grid_timestamps[1] - grid_timestamps[0]
         nframes, _, height, width = video.shape
-        factor = spatial_patch_size * spatial_merge_size
         min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
         total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
         max_pixels = max(
             min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
-            int(min_pixels * 1.05))
+            int(min_pixels * 1.05),
+        )
         max_pixels = ele.get("max_pixels", max_pixels)
         # min_pixels = (factor ** 2) * 52
         # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
@@ -323,11 +341,9 @@ def fetch_video(
         process_info.pop("type", None)
         process_info.pop("video", None)
         images = [
-            fetch_image({
-                "image": video_element,
-                **process_info
-            },
-                        size_factor=image_factor)
+            fetch_image(
+                {"image": video_element, **process_info}, size_factor=image_factor
+            )
             for video_element in ele["video"]
         ]
         nframes = ceil_by_factor(len(images), FRAME_FACTOR)
@@ -336,8 +352,7 @@ def fetch_video(
         return images, None, None, None
 
 
-def extract_vision_info(
-        conversations: list[dict] | list[list[dict]]) -> list[dict]:
+def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
     vision_infos = []
     if isinstance(conversations[0], dict):
         conversations = [conversations]
@@ -345,18 +360,21 @@ def extract_vision_info(
         for message in conversation:
             if isinstance(message["content"], list):
                 for ele in message["content"]:
-                    if ("image" in ele or "image_url" in ele or "video" in ele
-                            or ele["type"] in ("image", "image_url", "video")):
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele["type"] in ("image", "image_url", "video")
+                    ):
                         vision_infos.append(ele)
     return vision_infos
 
 
 def process_vision_info(
     conversations: list[dict] | list[list[dict]],
-) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]]
-           | None]:
+) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]:
     vision_infos = extract_vision_info(conversations)
-    ## Read images or videos
+    # Read images or videos
     image_inputs = []
     video_inputs = []
     for vision_info in vision_infos:
@@ -370,4 +388,4 @@ def process_vision_info(
         image_inputs = None
     if len(video_inputs) == 0:
         video_inputs = None
-    return image_inputs, video_inputs
+    return image_inputs, video_inputs

examples/offline_inference/qwen_2_5_omni/run_multiple_prompts.sh

@@ -0,0 +1,9 @@
+export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH
+python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --txt-prompts top100.txt

examples/offline_inference/qwen_2_5_omni/run_single_prompt.sh

@@ -1,9 +1,9 @@
 export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH
-export HF_ENDPOINT=https://hf-mirror.com
 python end2end.py --model Qwen/Qwen2.5-Omni-7B \
-                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." \
                                  --voice-type "m02" \
                                  --dit-ckpt none \
                                  --bigvgan-ckpt none \
                                  --output-wav output_audio \
-                                 --prompt_type text
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."

examples/offline_inference/qwen_2_5_omni/utils.py

@@ -1,30 +1,37 @@
 import tempfile
+from typing import Dict, List, Optional, Union
 from urllib.request import urlopen
-import librosa
-import soundfile as sf
-import resampy
-from typing import Dict, Optional
-import torch
-import requests
-import torchvision.io
 
-from typing import Union, List
+import librosa
+import requests
+import resampy
+import soundfile as sf
+import torch
+import torchvision.io
+from processing_omni import fetch_image, fetch_video
+
 from vllm.inputs import TextPrompt
 from vllm_omni.inputs.data import OmniTokensPrompt
-from processing_omni import fetch_image, fetch_video
+
+# Simple caches to avoid repeated heavy HF loads per prompt
+_PROCESSOR_CACHE: Dict[str, "AutoProcessor"] = {}
+_CONFIG_CACHE: Dict[str, "AutoConfig"] = {}
 
 
 def get_system_prompt():
-    
+
     return {
-        'role':
-        'system',
-        'content': [{
-            'type':
-            'text',
-            'text':
-            'You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.'
-        }]
+        "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."
+                ),
+            }
+        ],
     }
 
 
@@ -34,9 +41,7 @@ def resample_wav_to_16khz(input_filepath):
     if len(data.shape) > 1:
         data = data[:, 0]
     # resample to 16kHz
-    data_resampled = resampy.resample(data,
-                                      sr_orig=original_sample_rate,
-                                      sr_new=16000)
+    data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000)
     return data_resampled
 
 
@@ -55,8 +60,7 @@ def fetch_and_read_video(args, video_url: str, fps=2):
         total_duration = round(total_frames / video_fps, 3)
         nframes = int(total_frames / video_fps * fps)
 
-        frame_timestamps = total_duration * torch.arange(1,
-                                                         nframes + 1) / nframes
+        frame_timestamps = total_duration * torch.arange(1, nframes + 1) / nframes
         grid_timestamps = frame_timestamps[::2]
         second_per_grid = grid_timestamps[1] - grid_timestamps[0]
 
@@ -70,7 +74,8 @@ def fetch_and_read_video(args, video_url: str, fps=2):
 
     def read_video_with_transformers(video_file_name: Union[str, List[str]]):
         video, total_duration, nframes, second_per_grid = fetch_video(
-            {'video': video_file_name})
+            {"video": video_file_name}
+        )
         if total_duration is None and nframes is None:
             nframes = len(video)
             total_duration = 0.5 * nframes
@@ -89,7 +94,10 @@ def fetch_and_read_video(args, video_url: str, fps=2):
     if isinstance(video_url, str) and video_url.startswith("http"):
         with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
             resp = requests.get(video_url)
-            assert resp.status_code == requests.codes.ok, f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
+            assert resp.status_code == requests.codes.ok, (
+                f"Failed to fetch video from {video_url}, "
+                f"status_code:{resp.status_code}, resp:{resp}"
+            )
 
             temp_video_file.write(urlopen(video_url).read())
             temp_video_file_path = temp_video_file.name
@@ -106,104 +114,113 @@ def make_inputs_qwen2_omni(
     use_audio_in_video: Optional[bool] = False,
     tokenize: bool = False,
 ) -> Union[OmniTokensPrompt, TextPrompt]:
-    
-    from transformers import AutoConfig, AutoProcessor, AutoTokenizer
-    processor = AutoProcessor.from_pretrained(args.model)
-    tokenizer = AutoTokenizer.from_pretrained(args.model)
 
-    try:
-        config = AutoConfig.from_pretrained(args.model)
-        if 'Qwen2_5OmniModel' in config.architectures:
-            args.legacy_omni_video = False
+    from transformers import AutoConfig, AutoProcessor
+
+    # Cached processor/config to prevent per-prompt reloading and repeated warnings
+    if args.model not in _PROCESSOR_CACHE:
+        _PROCESSOR_CACHE[args.model] = AutoProcessor.from_pretrained(args.model)
+    processor = _PROCESSOR_CACHE[args.model]
+
+    config = _CONFIG_CACHE.get(args.model)
+    if config is None:
+        try:
+            config = AutoConfig.from_pretrained(args.model)
+        except Exception:
+            config = None
+        _CONFIG_CACHE[args.model] = config  # cache even if None to avoid retry storms
+
+    # Decide legacy flag only once based on config (default True if unknown)
+    if getattr(args, "legacy_omni_video", None) is None:
+        if config is not None and hasattr(config, "architectures"):
+            args.legacy_omni_video = not ("Qwen2_5OmniModel" in config.architectures)
         else:
             args.legacy_omni_video = True
-    except:
-        args.legacy_omni_video = True
 
     audios, images, videos = [], [], []
     for message in messages:
-        if not isinstance(message['content'], list):
-            message['content'] = [{
-                'type': 'text',
-                'text': message['content'],
-            }]
-        index, num_contents = 0, len(message['content'])
+        if not isinstance(message["content"], list):
+            message["content"] = [
+                {
+                    "type": "text",
+                    "text": message["content"],
+                }
+            ]
+        index, num_contents = 0, len(message["content"])
         while index < num_contents:
-            ele = message['content'][index]
-            if 'type' not in ele:
-                if 'text' in ele:
-                    ele['type'] = 'text'
-                elif 'audio' in ele:
-                    ele['type'] = 'audio'
-                elif 'audio_url' in ele:
-                    ele['type'] = 'audio_url'
-                elif 'image' in ele:
-                    ele['type'] = 'image'
-                elif 'image_url' in ele:
-                    ele['type'] = 'image_url'
-                elif 'video' in ele:
-                    ele['type'] = 'video'
-                elif 'video_url' in ele:
-                    ele['type'] = 'video_url'
+            ele = message["content"][index]
+            if "type" not in ele:
+                if "text" in ele:
+                    ele["type"] = "text"
+                elif "audio" in ele:
+                    ele["type"] = "audio"
+                elif "audio_url" in ele:
+                    ele["type"] = "audio_url"
+                elif "image" in ele:
+                    ele["type"] = "image"
+                elif "image_url" in ele:
+                    ele["type"] = "image_url"
+                elif "video" in ele:
+                    ele["type"] = "video"
+                elif "video_url" in ele:
+                    ele["type"] = "video_url"
                 else:
-                    raise ValueError(f'Unknown ele: {ele}')
+                    raise ValueError(f"Unknown ele: {ele}")
 
-            if ele['type'] == 'audio' or ele['type'] == 'audio_url':
-                if 'audio_url' in ele:
-                    audio_key = 'audio_url'
-                    with tempfile.NamedTemporaryFile(
-                            delete=True) as temp_audio_file:
+            if ele["type"] == "audio" or ele["type"] == "audio_url":
+                if "audio_url" in ele:
+                    audio_key = "audio_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_audio_file:
                         temp_audio_file.write(urlopen(ele[audio_key]).read())
                         temp_audio_file_path = temp_audio_file.name
-                        audios.append(
-                            resample_wav_to_16khz(temp_audio_file_path))
-                        ele['audio'] = temp_audio_file_path
-                elif 'audio' in ele:
-                    audio_key = 'audio'
+                        audios.append(resample_wav_to_16khz(temp_audio_file_path))
+                        ele["audio"] = temp_audio_file_path
+                elif "audio" in ele:
+                    audio_key = "audio"
                     audios.append(resample_wav_to_16khz(ele[audio_key]))
                 else:
-                    raise ValueError(f'Unknown ele {ele}')
-            elif use_audio_in_video and (ele['type'] == 'video'
-                                         or ele['type'] == 'video_url'):
+                    raise ValueError(f"Unknown ele {ele}")
+            elif use_audio_in_video and (
+                ele["type"] == "video" or ele["type"] == "video_url"
+            ):
                 # use video as audio as well
-                if 'video_url' in ele:
-                    audio_key = 'video_url'
-                    with tempfile.NamedTemporaryFile(
-                            delete=True) as temp_video_file:
+                if "video_url" in ele:
+                    audio_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
                         temp_video_file.write(urlopen(ele[audio_key]).read())
                         temp_video_file_path = temp_video_file.name
                         ele[audio_key] = temp_video_file_path
-                        audios.append(
-                            librosa.load(temp_video_file_path, sr=16000)[0])
-                        videos.append(
-                            fetch_and_read_video(args, temp_video_file_path))
-                        ele['video'] = temp_video_file_path
-                elif 'video' in ele:
-                    audio_key = 'video'
+                        audios.append(librosa.load(temp_video_file_path, sr=16000)[0])
+                        videos.append(fetch_and_read_video(args, temp_video_file_path))
+                        ele["video"] = temp_video_file_path
+                elif "video" in ele:
+                    audio_key = "video"
                     audios.append(librosa.load(ele[audio_key], sr=16000)[0])
                     videos.append(fetch_and_read_video(args, audio_key))
                 else:
                     raise ValueError("Unknown ele {}".format(ele))
                 # insert a audio after the video
-                message['content'].insert(index + 1, {
-                    "type": "audio",
-                    "audio": ele[audio_key],
-                })
+                message["content"].insert(
+                    index + 1,
+                    {
+                        "type": "audio",
+                        "audio": ele[audio_key],
+                    },
+                )
                 # no need to load the added audio again
                 index += 1
-            elif ele['type'] == 'video' or ele['type'] == 'video_url':
-                if 'video_url' in ele:
-                    video_key = 'video_url'
-                    with tempfile.NamedTemporaryFile(
-                            delete=True) as temp_video_file:
-                        temp_video_file.write(urlopen(ele['video_url']).read())
+            elif ele["type"] == "video" or ele["type"] == "video_url":
+                if "video_url" in ele:
+                    video_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele["video_url"]).read())
                         temp_video_file_path = temp_video_file.name
                         videos.append(fetch_and_read_video(args, temp_video_file))
-                        ele['video'] = temp_video_file_path
+                        ele["video"] = temp_video_file_path
                 else:
-                    video_key = 'video'
+                    video_key = "video"
                     videos.append(fetch_and_read_video(args, ele[video_key]))
-            elif ele['type'] == 'image' or ele['type'] == 'image_url':
+            elif ele["type"] == "image" or ele["type"] == "image_url":
                 images.append(fetch_image(ele))
 
             # move to the next content
@@ -249,11 +266,8 @@ def make_text_prompt(args, prompt):
         {
             "role": "user",
             "content": [
-                {
-                    "type": "text",
-                    "text": prompt
-                },
-            ]
+                {"type": "text", "text": prompt},
+            ],
         },
     ]
 
@@ -264,26 +278,29 @@ def make_text_prompt(args, prompt):
 def make_audio_in_video_v2_prompt(args):
     messages = [
         {
-            'role':
-            'system',
-            'content': [{
-                'type':
-                'text',
-                'text':
-                'You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.'
-            }]
-        },
-        {
-            "role":
-            "user",
+            "role": "system",
             "content": [
                 {
-                    "type":
-                    "video_url",
-                    "video_url":
-                    "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"
+                    "type": "text",
+                    "text": (
+                        "You are Qwen, a virtual human developed by the Qwen Team, "
+                        "Alibaba Group, capable of perceiving auditory and visual "
+                        "inputs, as well as generating text and speech."
+                    ),
+                }
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "video_url",
+                    "video_url": (
+                        "https://qianwen-res.oss-cn-beijing.aliyuncs.com/"
+                        "Qwen2.5-Omni/draw_small.mp4"
+                    ),
                 },
-            ]
+            ],
         },
     ]
     prompt = make_inputs_qwen2_omni(
@@ -295,11 +312,13 @@ def make_audio_in_video_v2_prompt(args):
     return prompt
 
 
-def make_omni_prompt(args, prompt = None) -> Union[OmniTokensPrompt, List[OmniTokensPrompt]]:
-    if args.prompt_type == 'text':
+def make_omni_prompt(
+    args, prompt=None
+) -> Union[OmniTokensPrompt, List[OmniTokensPrompt]]:
+    if args.prompt_type == "text":
         prompt = make_text_prompt(args, prompt)
-    elif args.prompt_type == 'audio-in-video-v2':
+    elif args.prompt_type == "audio-in-video-v2":
         prompt = make_audio_in_video_v2_prompt(args)
     else:
-        raise ValueError(f'Unsupported prompt type: {args.prompt_type}')
-    return prompt
+        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
+    return prompt

examples/omni/README.md

@@ -1,31 +0,0 @@
-# Omni Examples
-
-Examples showcasing multi-stage AR + DiT pipelines powered by vLLM-omni.
-
-## AR ➜ DiT with Diffusers Backend
-
-```bash
-# Optional: customise models
-export AR_MODEL=Qwen/Qwen3-0.6B
-export DIT_MODEL=./models/stable-diffusion-2-1/
-
-# Run with defaults (20 steps @ 512x512)
-python examples/omni/ar_dit_diffusers.py
-
-# Faster run
-python examples/omni/ar_dit_diffusers.py --steps 14 --height 384 --width 384 --guidance 4.5
-```
-
-Arguments:
-- `--ar-model`: Override AR stage model (default `Qwen/Qwen3-0.6B`).
-- `--dit-model`: Override DiT stage model (default env `DIT_MODEL`, else diffusers repo).
-- `--steps`: Number of diffusion steps (default 20).
-- `--guidance`: Guidance scale (default 5.0).
-- `--height` / `--width`: Output resolution (default 512).
-- `--seed`: Optional RNG seed.
-- `--prompt`, `--temperature`, `--max-tokens`: Control AR stage generation.
-- `--output`: Destination path for the generated image file.
-
-Generated images are saved next to the script as `omni_dit_output.png` by default.
-
-See the script for additional tweaks (scheduler, prompt chaining, etc.).

examples/omni/ar_dit_diffusers.py

@@ -1,175 +0,0 @@
-"""Run the AR → DiT (diffusers) pipeline using YAML configuration defaults."""
-
-from __future__ import annotations
-
-import argparse
-import os
-from pathlib import Path
-from typing import Dict, List
-
-import yaml
-
-from vllm_omni import (
-    OmniLLM,
-    DiTConfig,
-    create_ar_stage_config,
-    create_dit_stage_config,
-)
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="Run a two-stage AR → DiT pipeline with config-derived defaults."
-    )
-    parser.add_argument(
-        "--config",
-        type=Path,
-        default=Path(__file__).resolve().parent / "configs" / "ar_dit_local.yaml",
-        help="Path to a YAML configuration describing the pipeline stages.",
-    )
-    parser.add_argument(
-        "--prompt",
-        default="A scenic watercolor painting of a lighthouse at sunset",
-        help="Prompt passed to the AR stage.",
-    )
-    parser.add_argument(
-        "--negative-prompt",
-        default=None,
-        help="Optional negative prompt forwarded to the diffusion stage.",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=None,
-        help="Optional seed for deterministic diffusion sampling.",
-    )
-    parser.add_argument(
-        "--output",
-        type=Path,
-        default=Path("./omni_dit_output.png"),
-        help="Destination path for the generated image.",
-    )
-    return parser.parse_args()
-
-
-def _normalize_modalities(values: List[str] | None, default: List[str]) -> List[str]:
-    if not values:
-        return default
-    return list(values)
-
-
-def _load_stage_configs_from_yaml(config_path: Path):
-    with config_path.open("r", encoding="utf-8") as f:
-        config = yaml.safe_load(f) or {}
-
-    stages = config.get("stages", [])
-    if not stages:
-        raise ValueError(f"No stages defined in config: {config_path}")
-
-    stage_configs = []
-    for stage in stages:
-        stage_id = stage["stage_id"]
-        engine_type = stage["engine_type"]
-        default_stage_args = stage.get("default_stage_args")
-
-        if engine_type == "AR":
-            stage_configs.append(
-                create_ar_stage_config(
-                    stage_id=stage_id,
-                    model_path=stage["model_path"],
-                    input_modalities=_normalize_modalities(
-                        stage.get("input_modalities"), ["text"]
-                    ),
-                    output_modalities=_normalize_modalities(
-                        stage.get("output_modalities"), ["text"]
-                    ),
-                    default_stage_args=default_stage_args,
-                )
-            )
-        elif engine_type == "DiT":
-            dit_cfg_dict: Dict = dict(stage.get("dit_config", {}))
-            if "num_inference_steps" not in dit_cfg_dict:
-                raise ValueError(
-                    "DiT stage requires 'num_inference_steps' in dit_config"
-                )
-            dit_cfg = DiTConfig(**dit_cfg_dict)
-
-            stage_configs.append(
-                create_dit_stage_config(
-                    stage_id=stage_id,
-                    model_path=stage["model_path"],
-                    input_modalities=_normalize_modalities(
-                        stage.get("input_modalities"), ["text"]
-                    ),
-                    output_modalities=_normalize_modalities(
-                        stage.get("output_modalities"), ["image"]
-                    ),
-                    dit_config=dit_cfg,
-                    default_stage_args=default_stage_args,
-                )
-            )
-        else:
-            raise ValueError(f"Unsupported engine_type '{engine_type}' in config")
-
-    return stage_configs
-
-
-def _apply_env_overrides(stage_configs):
-    ar_override = os.environ.get("AR_MODEL")
-    dit_override = os.environ.get("DIT_MODEL")
-
-    for stage_config in stage_configs:
-        if ar_override and stage_config.engine_type == "AR":
-            stage_config.model_path = ar_override
-        if dit_override and stage_config.engine_type == "DiT":
-            stage_config.model_path = dit_override
-
-
-def main():
-    args = parse_args()
-
-    stage_configs = _load_stage_configs_from_yaml(args.config)
-    _apply_env_overrides(stage_configs)
-
-    omni = OmniLLM(stage_configs)
-
-    stage_overrides: Dict[int, Dict[str, object]] = {}
-    if args.seed is not None or args.negative_prompt:
-        for stage_config in stage_configs:
-            if stage_config.engine_type != "DiT":
-                continue
-            override = stage_overrides.setdefault(stage_config.stage_id, {})
-            if args.seed is not None:
-                override["seed"] = args.seed
-            if args.negative_prompt:
-                override["negative_prompt"] = args.negative_prompt
-
-    outputs = omni.generate(
-        prompt=args.prompt,
-        stage_overrides=stage_overrides if stage_overrides else None,
-    )
-
-    image = None
-    for request_output in outputs:
-        for completion in getattr(request_output, "outputs", []) or []:
-            candidate = getattr(completion, "image", None)
-            if candidate is not None:
-                image = candidate
-                break
-        if image is not None:
-            break
-
-    if image is None:
-        print("No image found in outputs.")
-        return
-
-    out_path = args.output.resolve()
-    try:
-        image.save(out_path)
-        print(f"Saved image to {out_path}")
-    except Exception:
-        print("Generated output is not a PIL.Image; skipping save.")
-
-
-if __name__ == "__main__":
-    main()

examples/omni/configs/ar_dit_local.yaml

@@ -1,28 +0,0 @@
-# Example config for running OmniLLM with local AR + DiT models.
-stages:
-  - stage_id: 0
-    engine_type: AR
-    model_path: Qwen/Qwen3-0.6B
-    input_modalities: [text]
-    output_modalities: [text]
-    default_stage_args:
-      max_tokens: 16
-      temperature: 0.7
-  - stage_id: 1
-    engine_type: DiT
-    model_path: stabilityai/stable-diffusion-2-1
-    input_modalities: [text]
-    output_modalities: [image]
-    default_stage_args:
-      height: 512
-      width: 512
-    dit_config:
-      num_inference_steps: 20
-      guidance_scale: 5.0
-      use_diffusers: true
-      diffusers_pipeline: auto
-      height: 512
-      width: 512
-      device_config:
-        device: mps
-        dtype: float16

pyproject.toml

@@ -103,6 +103,10 @@ multi_line_output = 3
 line_length = 88
 known_first_party = ["vllm_omni"]
 
+[tool.flake8]
+max-line-length = 88
+extend-ignore = ["E203", "W503"]
+
 [tool.mypy]
 python_version = "3.8"
 warn_return_any = true

requirements.txt

@@ -1,27 +1,4 @@
-# Core dependencies
-vllm>=0.10.2
-torch>=2.7
-transformers>=4.30.0
-numpy>=1.21.0
-pillow>=8.0.0
-opencv-python>=4.5.0
-PyYAML>=6.0
-
-# API and serving
-fastapi>=0.100.0
-uvicorn>=0.20.0
-pydantic>=2.0.0
-gradio>=3.0.0
-
-# Async and distributed computing
-asyncio-mqtt>=0.11.0
-ray>=2.0.0
-
-# Development dependencies
-pytest>=7.0.0
-pytest-asyncio>=0.21.0
-pytest-cov>=4.0.0
-black>=23.0.0
-isort>=5.12.0
-flake8>=6.0.0
-mypy>=1.0.0
+omegaconf==2.3.0
+decord==0.6.0
+librosa==0.11.0
+resampy==0.4.3

scripts/README.md

@@ -106,4 +106,3 @@ python scripts/download_models.py
 - `server.log`: Server startup and runtime logs
 - `api_client_test.log`: API client test results
 - `simple_usage_test.log`: Simple usage test results
-

scripts/test_serving.sh

@@ -99,15 +99,15 @@ print('✅ All imports successful')
 start_server() {
     log_info "Starting vLLM-omni server on port $PORT..."
     log_info "Command: vllm serve $MODEL_PATH --omni --port $PORT"
-    
+
     # Start server in background
     vllm serve "$MODEL_PATH" --omni --port "$PORT" > server.log 2>&1 &
     SERVER_PID=$!
-    
+
     log_info "Server started with PID: $SERVER_PID"
     log_info "Waiting $SERVER_STARTUP_WAIT seconds for server to initialize..."
     sleep $SERVER_STARTUP_WAIT
-    
+
     # Check if server is still running
     if ! kill -0 $SERVER_PID 2>/dev/null; then
         log_error "Server failed to start"
@@ -115,7 +115,7 @@ start_server() {
         cat server.log
         exit 1
     fi
-    
+
     log_success "Server appears to be running"
 }
 
@@ -146,7 +146,7 @@ test_info() {
     local response=$(curl -s -w "%{http_code}" http://$HOST:$PORT/info 2>/dev/null || echo "000")
     local http_code="${response: -3}"
     local body="${response%???}"
-    
+
     if [ "$http_code" = "200" ]; then
         log_success "Info endpoint working"
         echo "$body" | python -m json.tool 2>/dev/null || log_warning "Info response not valid JSON"
@@ -159,13 +159,13 @@ test_info() {
 # Test text generation
 test_generation() {
     log_info "Testing text generation..."
-    
+
     # Test 1: Simple generation
     local response=$(curl -s -X POST http://$HOST:$PORT/generate \
         -H "Content-Type: application/json" \
         -d '{"prompts": ["Test the server functionality"], "max_tokens": 20, "temperature": 0.7}' \
         2>/dev/null || echo "{}")
-    
+
     if echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); exit(0 if 'outputs' in data and len(data['outputs']) > 0 else 1)" 2>/dev/null; then
         log_success "Text generation test passed"
         echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); print('Generated text:', data['outputs'][0]['text'][:100] + '...' if len(data['outputs'][0]['text']) > 100 else data['outputs'][0]['text'])" 2>/dev/null
@@ -271,7 +271,7 @@ test_simple_usage() {
 # Performance test
 test_performance() {
     log_info "Running performance test..."
-    
+
     local start_time=$(date +%s)
     local response=$(curl -s -X POST http://$HOST:$PORT/generate \
         -H "Content-Type: application/json" \
@@ -279,7 +279,7 @@ test_performance() {
         2>/dev/null)
     local end_time=$(date +%s)
     local duration=$((end_time - start_time))
-    
+
     if echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); exit(0 if 'outputs' in data else 1)" 2>/dev/null; then
         log_success "Performance test passed (${duration}s response time)"
     else
@@ -295,16 +295,16 @@ run_tests() {
     log_info "Port: $PORT"
     log_info "Host: $HOST"
     echo "=========================================="
-    
+
     # Run all tests
     check_model
     check_environment
     test_imports
     start_server
-    
+
     # Wait a bit more for server to be fully ready
     sleep 5
-    
+
     test_health
     test_info
     test_generation
@@ -348,13 +348,13 @@ main() {
         show_usage
         exit 0
     fi
-    
+
     # Change to script directory
     cd "$(dirname "$0")/.."
-    
+
     # Create logs directory if it doesn't exist
     mkdir -p logs
-    
+
     # Run tests
     run_tests
 }

tests/__init__.py

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

tests/conftest.py

@@ -1,82 +0,0 @@
-"""
-Shared fixtures and configuration for vLLM-omni tests.
-"""
-
-import pytest
-import torch
-from unittest.mock import Mock
-from vllm_omni.config import OmniStageConfig, DiTConfig, DiTCacheConfig, create_ar_stage_config, create_dit_stage_config
-
-
-@pytest.fixture(scope="session")
-def device():
-    """Get available device for testing."""
-    return "cuda" if torch.cuda.is_available() else "cpu"
-
-
-@pytest.fixture
-def sample_ar_stage_config():
-    """Sample AR stage configuration for testing."""
-    return create_ar_stage_config(
-        stage_id=0,
-        model_path="test-ar-model",
-        input_modalities=["text"],
-        output_modalities=["text"]
-    )
-
-
-@pytest.fixture
-def sample_dit_stage_config():
-    """Sample DiT stage configuration for testing."""
-    dit_config = DiTConfig(
-        model_type="dit",
-        scheduler_type="ddpm",
-        num_inference_steps=10,
-        guidance_scale=7.5
-    )
-    
-    return create_dit_stage_config(
-        stage_id=1,
-        model_path="test-dit-model",
-        input_modalities=["text"],
-        output_modalities=["image"],
-        dit_config=dit_config
-    )
-
-
-@pytest.fixture
-def sample_stage_configs(sample_ar_stage_config, sample_dit_stage_config):
-    """Sample stage configurations for testing."""
-    return [sample_ar_stage_config, sample_dit_stage_config]
-
-
-@pytest.fixture
-def mock_vllm_config():
-    """Mock vLLM configuration."""
-    config = Mock()
-    config.model = "test-model"
-    config.tensor_parallel_size = 1
-    config.pipeline_parallel_size = 1
-    return config
-
-
-@pytest.fixture
-def mock_dit_cache_config():
-    """Mock DiT cache configuration."""
-    from vllm_omni.config import DiTCacheTensor
-    
-    cache_tensors = [
-        DiTCacheTensor(
-            name="test_tensor",
-            shape=[1, 512, 512],
-            dtype="float32",
-            persistent=True
-        )
-    ]
-    
-    return DiTCacheConfig(
-        cache_tensors=cache_tensors,
-        max_cache_size=1024 * 1024 * 1024,  # 1GB
-        cache_strategy="fifo",
-        enable_optimization=True
-    )

tests/unit/test_config.py

@@ -1,179 +0,0 @@
-"""
-Unit tests for configuration modules.
-"""
-
-import pytest
-from vllm_omni.config import (
-    OmniStageConfig,
-    DiTConfig,
-    DiTCacheConfig,
-    DiTCacheTensor,
-    create_ar_stage_config,
-    create_dit_stage_config,
-)
-
-
-class TestOmniStageConfig:
-    def test_ar_stage_config_creation(self):
-        """Test AR stage configuration creation."""
-        config = OmniStageConfig(
-            stage_id=0,
-            engine_type="AR",
-            model_path="test-model",
-            input_modalities=["text"],
-            output_modalities=["text"]
-        )
-        assert config.engine_type == "AR"
-        assert config.stage_id == 0
-        assert config.model_path == "test-model"
-        assert config.input_modalities == ["text"]
-        assert config.output_modalities == ["text"]
-    
-    def test_dit_stage_config_creation(self):
-        """Test DiT stage configuration creation."""
-        dit_config = DiTConfig(
-            model_type="dit",
-            scheduler_type="ddpm",
-            num_inference_steps=50
-        )
-        config = OmniStageConfig(
-            stage_id=1,
-            engine_type="DiT",
-            model_path="test-dit-model",
-            input_modalities=["text"],
-            output_modalities=["image"],
-            dit_config=dit_config
-        )
-        assert config.engine_type == "DiT"
-        assert config.dit_config is not None
-        assert config.dit_config.model_type == "dit"
-    
-    def test_invalid_engine_type(self):
-        """Test validation of engine type."""
-        with pytest.raises(ValueError):
-            OmniStageConfig(
-                stage_id=0,
-                engine_type="INVALID",
-                model_path="test-model",
-                input_modalities=["text"],
-                output_modalities=["text"]
-            )
-    
-    def test_empty_modalities(self):
-        """Test validation of empty modalities."""
-        with pytest.raises(ValueError):
-            OmniStageConfig(
-                stage_id=0,
-                engine_type="AR",
-                model_path="test-model",
-                input_modalities=[],
-                output_modalities=["text"]
-            )
-        
-        with pytest.raises(ValueError):
-            OmniStageConfig(
-                stage_id=0,
-                engine_type="AR",
-                model_path="test-model",
-                input_modalities=["text"],
-                output_modalities=[]
-            )
-    
-    def test_dit_requires_config(self):
-        """Test that DiT engine requires dit_config."""
-        with pytest.raises(ValueError):
-            OmniStageConfig(
-                stage_id=1,
-                engine_type="DiT",
-                model_path="test-dit-model",
-                input_modalities=["text"],
-                output_modalities=["image"]
-            )
-
-
-class TestDiTConfig:
-    def test_dit_config_defaults(self):
-        """Test DiT configuration with defaults."""
-        config = DiTConfig(
-            model_type="dit",
-            scheduler_type="ddpm",
-            num_inference_steps=50
-        )
-        assert config.use_diffusers is False
-        assert config.diffusers_pipeline is None
-        assert config.guidance_scale == 7.5
-        assert config.height == 512
-        assert config.width == 512
-    
-    def test_diffusers_config(self):
-        """Test DiT configuration with diffusers."""
-        config = DiTConfig(
-            model_type="dit",
-            scheduler_type="ddpm",
-            num_inference_steps=50,
-            use_diffusers=True,
-            diffusers_pipeline="stable-diffusion"
-        )
-        assert config.use_diffusers is True
-        assert config.diffusers_pipeline == "stable-diffusion"
-
-
-class TestDiTCacheConfig:
-    def test_cache_config_creation(self):
-        """Test cache configuration creation."""
-        cache_tensors = [
-            DiTCacheTensor(
-                name="test_tensor",
-                shape=[1, 512, 512],
-                dtype="float32",
-                persistent=True
-            )
-        ]
-        
-        config = DiTCacheConfig(
-            cache_tensors=cache_tensors,
-            max_cache_size=1024 * 1024 * 1024,
-            cache_strategy="fifo"
-        )
-        
-        assert len(config.cache_tensors) == 1
-        assert config.max_cache_size == 1024 * 1024 * 1024
-        assert config.cache_strategy == "fifo"
-        assert config.enable_optimization is True
-
-
-class TestHelperFunctions:
-    def test_create_ar_stage_config(self):
-        """Test create_ar_stage_config helper function."""
-        config = create_ar_stage_config(
-            stage_id=0,
-            model_path="test-model"
-        )
-        
-        assert config.stage_id == 0
-        assert config.engine_type == "AR"
-        assert config.model_path == "test-model"
-        assert config.input_modalities == ["text"]
-        assert config.output_modalities == ["text"]
-    
-    def test_create_dit_stage_config(self):
-        """Test create_dit_stage_config helper function."""
-        dit_config = DiTConfig(
-            model_type="dit",
-            scheduler_type="ddpm",
-            num_inference_steps=50
-        )
-        
-        config = create_dit_stage_config(
-            stage_id=1,
-            model_path="test-dit-model",
-            dit_config=dit_config
-        )
-        
-        assert config.stage_id == 1
-        assert config.engine_type == "DiT"
-        assert config.model_path == "test-dit-model"
-        assert config.input_modalities == ["text"]
-        assert config.output_modalities == ["image"]
-        assert config.dit_config is not None
-

vllm_omni/__init__.py

@@ -1,45 +1,29 @@
 """
-vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures.
+vLLM-omni: Multi-modality models inference and serving with
+non-autoregressive structures.
 
-This package extends vLLM beyond traditional text-based, autoregressive generation
-to support multi-modality models with non-autoregressive structures and non-textual outputs.
+This package extends vLLM beyond traditional text-based, autoregressive
+generation to support multi-modality models with non-autoregressive
+structures and non-textual outputs.
 
 Architecture:
-- 🟡 Modified: vLLM components modified for multimodal support  
-- 🔴 Added: New components for multimodal and non-autoregressive processing
+- 🟡 Modified: vLLM components modified for multimodal support
+- 🔴 Added: New components for multimodal and non-autoregressive
+  processing
 """
 
-__version__ = "0.1.0"
-__author__ = "vLLM-omni Team"
-__email__ = "hsliuustc@gmail.com"
+from . import patch  # noqa: F401
+from .config import OmniModelConfig
 
 # Main entry points
-from .entrypoints.omni_llm import OmniLLM, AsyncOmniLLM
-from .config import (
-    OmniStageConfig,
-    DiTConfig,
-    DiTCacheConfig,
-    create_ar_stage_config,
-    create_dit_stage_config,
-)
+from .entrypoints.omni_llm import OmniLLM
 
 __all__ = [
-    # Version info
-    "__version__",
-    "__author__", 
-    "__email__",
-    
     # Main components
     "OmniLLM",
-    "AsyncOmniLLM",
-    
+    # "AsyncOmniLLM", # TODO: add AsyncOmniLLM back when it is implemented
     # Configuration
-    "OmniStageConfig",
-    "DiTConfig",
-    "DiTCacheConfig",
-    "create_ar_stage_config",
-    "create_dit_stage_config",
-    
+    "OmniModelConfig",
     # All other components are available through their respective modules
     # processors.*, schedulers.*, executors.*, etc.
-]
+]

vllm_omni/config/__init__.py

@@ -1,50 +1,9 @@
 """
 Configuration module for vLLM-omni.
 """
-from typing import Optional
-from pydantic.dataclasses import dataclass
-from pydantic import ConfigDict
-
-from vllm.config import ModelConfig
-from vllm.config import config
-
-import vllm_omni.model_executor.models as me_models
-
-from .stage_config import (
-    OmniStageConfig,
-    DiTConfig,
-    DiTCacheConfig,
-    DiTCacheTensor,
-    create_ar_stage_config,
-    create_dit_stage_config,
-)
-
-
-@config
-@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
-class OmniModelConfig(ModelConfig):
-    """Configuration for Omni models, extending the base ModelConfig."""
-    
-    stage_id: int = 0
-    model_stage: str = "thinker"
-    model_arch: str = "Qwen2_5OmniForConditionalGeneration"
-    engine_output_type: Optional[str] = None
-    
-    @property
-    def registry(self):
-        return me_models.OmniModelRegistry
-    
-    @property
-    def architectures(self) -> list[str]:
-        return [self.model_arch]
 
+from vllm_omni.config.model import OmniModelConfig
 
 __all__ = [
     "OmniModelConfig",
-    "OmniStageConfig",
-    "DiTConfig", 
-    "DiTCacheConfig",
-    "DiTCacheTensor",
-    "create_ar_stage_config",
-    "create_dit_stage_config",
-]
+]

vllm_omni/config/model.py

@@ -0,0 +1,26 @@
+from typing import Optional
+
+from pydantic import ConfigDict
+from pydantic.dataclasses import dataclass
+
+import vllm_omni.model_executor.models as me_models
+from vllm.config import ModelConfig, config
+
+
+@config
+@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
+class OmniModelConfig(ModelConfig):
+    """Configuration for Omni models, extending the base ModelConfig."""
+
+    stage_id: int = 0
+    model_stage: str = "thinker"
+    model_arch: str = "Qwen2_5OmniForConditionalGeneration"
+    engine_output_type: Optional[str] = None
+
+    @property
+    def registry(self):
+        return me_models.OmniModelRegistry
+
+    @property
+    def architectures(self) -> list[str]:
+        return [self.model_arch]

vllm_omni/config/stage_config.py

@@ -1,166 +0,0 @@
-"""
-Stage configuration for vLLM-omni multi-stage processing.
-"""
-
-from dataclasses import dataclass
-from typing import List, Optional, Type, Literal, Any, Dict
-from vllm.config import (
-    VllmConfig,
-    DeviceConfig,
-    LoadConfig,
-    ModelConfig,
-    CompilationConfig,
-)
-from vllm.executor.executor_base import ExecutorBase as Executor
-
-@dataclass
-class DiTCacheTensor:
-    """Configuration for DiT cache tensors."""
-    name: str
-    shape: List[int]
-    dtype: str = "float32"
-    persistent: bool = True
-
-
-@dataclass
-class DiTCacheConfig:
-    """Configuration for DiT caching system."""
-    cache_tensors: List[DiTCacheTensor]
-    max_cache_size: int = 1024 * 1024 * 1024  # 1GB
-    cache_strategy: str = "fifo"  # fifo, lru, lfu
-    enable_optimization: bool = True
-    cache_compression: bool = False
-
-
-@dataclass
-class DiTConfig:
-    """Configuration for DiT (Diffusion Transformer) stages."""
-
-    num_inference_steps: int
-    """Number of diffusion inference steps."""
-
-    model_config: Optional[ModelConfig] = None
-    """Model configuration."""
-    scheduler_config: Optional[Any] = None
-    """Scheduler configuration."""
-    device_config: Optional[DeviceConfig] = None
-    """Device configuration."""
-    load_config: Optional[LoadConfig] = None
-    """Model loading configuration."""
-    compilation_config: Optional[CompilationConfig] = None
-    """Compilation configuration."""
-    dit_cache_config: Optional[DiTCacheConfig] = None
-    """DiT cache configuration."""
-    guidance_scale: float = 7.5
-    use_diffusers: bool = False
-    diffusers_pipeline: Optional[str] = None
-    height: int = 512
-    width: int = 512
-    batch_size: int = 1
-
-
-@dataclass
-class OmniStageConfig:
-    """Configuration for a processing stage in vLLM-omni."""
-
-    stage_id: int
-    engine_type: Literal["AR", "DiT"]
-    model_path: str
-    input_modalities: List[str]
-    output_modalities: List[str]
-    vllm_config: Optional[VllmConfig] = None
-    dit_config: Optional[DiTConfig] = None
-    executor_class: Type[Executor] = None  # Will be set based on engine_type
-    default_stage_args: Optional[Dict[str, Any]] = None
-    stage_output: Optional[Any] = None
-    
-    def __post_init__(self):
-        """Validate configuration after initialization."""
-        if self.engine_type not in ["AR", "DiT"]:
-            raise ValueError(f"Invalid engine_type: {self.engine_type}. Must be 'AR' or 'DiT'")
-        
-        if self.engine_type == "DiT" and self.dit_config is None:
-            raise ValueError("DiT engine requires dit_config")
-        
-        if not self.input_modalities:
-            raise ValueError("input_modalities cannot be empty")
-        
-        if not self.output_modalities:
-            raise ValueError("output_modalities cannot be empty")
-    
-    def get_executor_class(self) -> Type[Executor]:
-        """Get the appropriate executor class for this stage."""
-        if self.executor_class is not None:
-            return self.executor_class
-        
-        if self.engine_type == "AR":
-            from vllm.executor.uniproc_executor import UniProcExecutor
-            return UniProcExecutor
-        elif self.engine_type == "DiT":
-            if self.dit_config and self.dit_config.use_diffusers:
-                from vllm_omni.executor.diffusers_executor import DiffusersPipelineExecutor
-                return DiffusersPipelineExecutor
-            else:
-                from vllm.executor.uniproc_executor import UniProcExecutor
-                return UniProcExecutor
-        
-        raise ValueError(f"No executor class available for engine_type: {self.engine_type}")
-
-
-def create_ar_stage_config(
-    stage_id: int,
-    model_path: str,
-    input_modalities: List[str] = None,
-    output_modalities: List[str] = None,
-    vllm_config: Optional[VllmConfig] = None,
-    executor_class: Optional[Type[Executor]] = None,
-    default_stage_args: Optional[Dict[str, Any]] = None,
-) -> OmniStageConfig:
-    """Create a configuration for an AR (Autoregressive) stage."""
-    if input_modalities is None:
-        input_modalities = ["text"]
-    if output_modalities is None:
-        output_modalities = ["text"]
-    
-    # For now, we'll create minimal configs
-    # vllm_config and executor_class will be handled by the LLM class
-    return OmniStageConfig(
-        stage_id=stage_id,
-        engine_type="AR",
-        model_path=model_path,
-        input_modalities=input_modalities,
-        output_modalities=output_modalities,
-        vllm_config=vllm_config,
-        executor_class=executor_class,
-        default_stage_args=default_stage_args,
-    )
-
-
-def create_dit_stage_config(
-    stage_id: int,
-    model_path: str,
-    input_modalities: Optional[List[str]] = None,
-    output_modalities: Optional[List[str]] = None,
-    dit_config: Optional[DiTConfig] = None,
-    executor_class: Optional[Type[Executor]] = None,
-    default_stage_args: Optional[Dict[str, Any]] = None,
-) -> OmniStageConfig:
-    """Create a configuration for a DiT (Diffusion Transformer) stage."""
-    if input_modalities is None:
-        input_modalities = ["text"]
-    if output_modalities is None:
-        output_modalities = ["image"]
-    
-    if dit_config is None:
-        dit_config = DiTConfig(num_inference_steps=50)
-
-    return OmniStageConfig(
-        stage_id=stage_id,
-        engine_type="DiT",
-        model_path=model_path,
-        input_modalities=input_modalities,
-        output_modalities=output_modalities,
-        dit_config=dit_config,
-        executor_class=executor_class,
-        default_stage_args=default_stage_args,
-    )

vllm_omni/core/dit_cache_manager.py

@@ -5,174 +5,9 @@ This module provides caching functionality for DiT (Diffusion Transformer) model
 to optimize inference performance and memory usage.
 """
 
-import time
-import torch
-from typing import Dict, List, Optional, Any
-from dataclasses import dataclass
-from ..config import DiTCacheConfig, DiTCacheTensor
-
 
 class DiTCacheManager:
     """Manages DiT-specific caching for optimized inference."""
-    
-    def __init__(self, config: DiTCacheConfig):
-        self.config = config
-        self.cache_tensors: Dict[str, DiTCacheTensor] = {}
-        self.cache_groups: List[DiTCacheTensor] = config.cache_tensors
-        self.max_cache_size = config.max_cache_size
-        self.current_cache_size = 0
-        self.cache_hits = 0
-        self.cache_misses = 0
-        self.cache_timeout = 3600  # 1 hour default timeout
-    
-    def allocate_cache(self, request_id: str, size: int) -> torch.Tensor:
-        """Allocate cache for a specific request."""
-        # Check if we have enough space
-        if self.current_cache_size + size > self.max_cache_size:
-            self._evict_cache()
-        
-        # Create tensor
-        tensor = torch.zeros(size, dtype=torch.float32)
-        
-        # Store in cache
-        cache_tensor = DiTCacheTensor(
-            name=request_id,
-            tensor=tensor,
-            timestamp=time.time(),
-            access_count=0,
-            size_bytes=size * 4  # Assuming float32
-        )
-        
-        self.cache_tensors[request_id] = cache_tensor
-        self.current_cache_size += cache_tensor.size_bytes
-        
-        return tensor
-    
-    def get_cache(self, request_id: str) -> Optional[torch.Tensor]:
-        """Get cached tensor for a request."""
-        if request_id in self.cache_tensors:
-            cache_tensor = self.cache_tensors[request_id]
-            cache_tensor.access_count += 1
-            self.cache_hits += 1
-            return cache_tensor.tensor
-        else:
-            self.cache_misses += 1
-            return None
-    
-    def store_cache(self, request_id: str, tensor: torch.Tensor) -> None:
-        """Store a tensor in the cache."""
-        if request_id in self.cache_tensors:
-            # Update existing cache
-            old_tensor = self.cache_tensors[request_id]
-            self.current_cache_size -= old_tensor.size_bytes
-            
-            new_size_bytes = tensor.numel() * 4  # Assuming float32
-            self.current_cache_size += new_size_bytes
-            
-            self.cache_tensors[request_id] = DiTCacheTensor(
-                name=request_id,
-                tensor=tensor,
-                timestamp=time.time(),
-                access_count=old_tensor.access_count,
-                size_bytes=new_size_bytes
-            )
-        else:
-            # Create new cache entry
-            new_size_bytes = tensor.numel() * 4
-            if self.current_cache_size + new_size_bytes > self.max_cache_size:
-                self._evict_cache()
-            
-            self.cache_tensors[request_id] = DiTCacheTensor(
-                name=request_id,
-                tensor=tensor,
-                timestamp=time.time(),
-                access_count=0,
-                size_bytes=new_size_bytes
-            )
-            self.current_cache_size += new_size_bytes
-    
-    def release_cache(self, request_id: str) -> None:
-        """Release cache for a request."""
-        if request_id in self.cache_tensors:
-            cache_tensor = self.cache_tensors[request_id]
-            self.current_cache_size -= cache_tensor.size_bytes
-            del self.cache_tensors[request_id]
-    
-    def clear_expired_cache(self) -> None:
-        """Clear expired cache entries."""
-        current_time = time.time()
-        expired_keys = []
-        
-        for request_id, cache_tensor in self.cache_tensors.items():
-            if current_time - cache_tensor.timestamp > self.cache_timeout:
-                expired_keys.append(request_id)
-        
-        for request_id in expired_keys:
-            self.release_cache(request_id)
-    
-    def _evict_cache(self) -> None:
-        """Evict cache entries based on strategy."""
-        if self.config.cache_strategy == "fifo":
-            self._evict_fifo()
-        elif self.config.cache_strategy == "lru":
-            self._evict_lru()
-        elif self.config.cache_strategy == "lfu":
-            self._evict_lfu()
-        else:
-            self._evict_fifo()  # Default to FIFO
-    
-    def _evict_fifo(self) -> None:
-        """Evict cache entries using FIFO strategy."""
-        if not self.cache_tensors:
-            return
-        
-        # Find oldest entry
-        oldest_key = min(self.cache_tensors.keys(), 
-                        key=lambda k: self.cache_tensors[k].timestamp)
-        self.release_cache(oldest_key)
-    
-    def _evict_lru(self) -> None:
-        """Evict cache entries using LRU strategy."""
-        if not self.cache_tensors:
-            return
-        
-        # Find least recently used entry
-        lru_key = min(self.cache_tensors.keys(),
-                     key=lambda k: self.cache_tensors[k].access_count)
-        self.release_cache(lru_key)
-    
-    def _evict_lfu(self) -> None:
-        """Evict cache entries using LFU strategy."""
-        if not self.cache_tensors:
-            return
-        
-        # Find least frequently used entry
-        lfu_key = min(self.cache_tensors.keys(),
-                     key=lambda k: self.cache_tensors[k].access_count)
-        self.release_cache(lfu_key)
-    
-    def get_statistics(self) -> Dict[str, Any]:
-        """Get cache statistics."""
-        total_requests = self.cache_hits + self.cache_misses
-        hit_rate = self.cache_hits / total_requests if total_requests > 0 else 0
-        
-        return {
-            "cache_size": self.current_cache_size,
-            "max_cache_size": self.max_cache_size,
-            "cache_utilization": self.current_cache_size / self.max_cache_size,
-            "cache_hits": self.cache_hits,
-            "cache_misses": self.cache_misses,
-            "hit_rate": hit_rate,
-            "num_cached_tensors": len(self.cache_tensors)
-        }
-    
-    def clear_all_cache(self) -> None:
-        """Clear all cache entries."""
-        self.cache_tensors.clear()
-        self.current_cache_size = 0
-        self.cache_hits = 0
-        self.cache_misses = 0
-    
-    def set_cache_timeout(self, timeout: float) -> None:
-        """Set cache timeout in seconds."""
-        self.cache_timeout = timeout
+
+    def __init__(self, config):
+        pass

vllm_omni/core/sched/__init__.py

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

vllm_omni/core/sched/diffusion_scheduler.py

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

vllm_omni/core/sched/output.py

@@ -0,0 +1,33 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.request import Request
+from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
+
+
+@dataclass
+class OmniNewRequestData(NewRequestData):
+    # Optional serialized prompt embeddings
+    prompt_embeds: Optional[PromptEmbedsPayload] = None
+    # Optional serialized additional information
+    additional_information: Optional[AdditionalInformationPayload] = None
+
+    @classmethod
+    def from_request(
+        cls,
+        request: Request,
+        block_ids: tuple[list[int], ...],
+    ) -> "OmniNewRequestData":
+        return cls(
+            req_id=request.request_id,
+            prompt_token_ids=request.prompt_token_ids,
+            mm_features=request.mm_features,
+            sampling_params=request.sampling_params,
+            pooling_params=request.pooling_params,
+            block_ids=block_ids,
+            num_computed_tokens=request.num_computed_tokens,
+            lora_request=request.lora_request,
+            prompt_embeds=request.prompt_embeds,
+            additional_information=request.additional_information,
+        )

vllm_omni/core/sched/scheduler.py

@@ -1,40 +1,60 @@
+from __future__ import annotations
 
-from vllm_omni.request import OmniRequest
-from typing import List
-from threading import Lock, Condition
-from vllm_omni.config import OmniConfig
-
-from vllm.v1.kv_cache_interface import KVCacheConfig
-from vllm.v1.core.sched import SchedulerInterface
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.v1.structured_output import StructuredOutputManager
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 
 
-class OmniScheduler(SchedulerInterface):
+class OmniScheduler(VLLMScheduler):
     """
     OmniScheduler: Scheduler for vLLM-omni multimodal processing.
 
-    This scheduler extends vLLM's scheduler to support multimodal and non-autoregressive
-    processing with additional fields and methods specific to vLLM-omni.
+    This scheduler extends vLLM's scheduler to support multimodal and
+    non-autoregressive processing with additional fields and methods
+    specific to vLLM-omni.
     """
-    def __init__(self,
-        omni_config: OmniConfig,
-        kv_cache_config: KVCacheConfig,
-        structured_output_manager: StructuredOutputManager,
-        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
-        include_finished_set: bool = False,
-        log_stats: bool = False,
-    ):
-        super().__init__(
-            vllm_config=omni_config.vllm_config,
-            kv_cache_config=kv_cache_config,
-            multimodal_registry=mm_registry,
-            structured_output_manager=structured_output_manager,
-            include_finished_set=include_finished_set,
-            log_stats=log_stats,
-        )
-        self.omni_config = omni_config
-            
-    def schedule(self, requests: List[OmniRequest]) -> List[OmniRequest]:
-        # TODO: Implement scheduling logic
-        pass
+
+    # Ensure scheduled_new_reqs carry omni-specific payloads
+    # (e.g., additional_information)
+    def schedule(self) -> SchedulerOutput:  # type: ignore[override]
+        scheduler_output = super().schedule()
+        try:
+            # Late import to avoid circulars in some launch modes
+            from .output import OmniNewRequestData
+
+            # Rewrap base NewRequestData entries with OmniNewRequestData,
+            # enriching with request-level payloads
+            new_list = []
+            for nr in scheduler_output.scheduled_new_reqs:
+                req_id = getattr(nr, "req_id", None)
+                request = self.requests.get(req_id) if req_id else None
+                # Build omni entry preserving all base fields
+                omni_nr = OmniNewRequestData(
+                    req_id=nr.req_id,
+                    prompt_token_ids=nr.prompt_token_ids,
+                    mm_features=nr.mm_features,
+                    sampling_params=nr.sampling_params,
+                    pooling_params=nr.pooling_params,
+                    block_ids=nr.block_ids,
+                    num_computed_tokens=nr.num_computed_tokens,
+                    lora_request=nr.lora_request,
+                    # Enrich with omni payloads from the live request object
+                    prompt_embeds=(
+                        getattr(request, "prompt_embeds", None) if request else None
+                    ),
+                    additional_information=(
+                        getattr(request, "additional_information", None)
+                        if request
+                        else None
+                    ),
+                )
+                new_list.append(omni_nr)
+
+            scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
+        except Exception:
+            # If anything goes wrong, leave the original output unchanged
+            init_logger(__name__).exception(
+                "Failed to wrap scheduled_new_reqs with OmniNewRequestData"
+            )
+
+        return scheduler_output

vllm_omni/engine/__init__.py

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

vllm_omni/engine/arg_utils.py

@@ -1,8 +1,7 @@
-from typing import Optional
 from dataclasses import dataclass
-from vllm.engine.arg_utils import EngineArgs
-from vllm.utils import FlexibleArgumentParser
+from typing import Optional
 
+from vllm.engine.arg_utils import EngineArgs
 from vllm_omni.config import OmniModelConfig
 
 
@@ -13,39 +12,22 @@ class OmniEngineArgs(EngineArgs):
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: Optional[str] = None
 
-    @staticmethod
-    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
-        """Shared CLI arguments for vLLM engine."""
-        parser.add_argument(
-            "--engine-output-type",
-            type=str,
-            default=EngineArgs.engine_output_type,
-            help=(
-                "Declare EngineCoreOutput.output_type (e.g., 'text', 'image', "
-                "'text+image', 'latent'). This will be written into "
-                "model_config.engine_output_type for schedulers to use."
-            ),
-        )
-        parser.add_argument("--model-stage", type=str, default=OmniEngineArgs.model_stage, 
-        help="Declare model stage (e.g., 'thinker', 'talker', 'token2wav'). This will be written into model_config.model_stage for schedulers to use.")
-        return parser
-
     def create_model_config(self) -> OmniModelConfig:
         # First, get the base ModelConfig from the parent class
         base_config = super().create_model_config()
-        
+
         # Create OmniModelConfig by copying all base config attributes
         # and adding the new omni-specific fields
         config_dict = base_config.__dict__.copy()
-        
+
         # Add the new omni-specific fields
-        config_dict['stage_id'] = self.stage_id
-        config_dict['model_stage'] = self.model_stage
-        config_dict['model_arch'] = self.model_arch
-        config_dict['engine_output_type'] = self.engine_output_type
-        
+        config_dict["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
+
+        return omni_config

vllm_omni/engine/diffusion_engine.py

@@ -1,133 +0,0 @@
-"""Diffusion engine adapters for diffusers-backed DiT stages."""
-
-from __future__ import annotations
-
-from typing import Any, Optional
-
-from ..config import DiTConfig
-from ..worker.gpu_diffusion_worker import (
-    DiffusionGPUWorker,
-    DiffusionRunnerOutput,
-)
-
-
-class DiffusionEngine:
-    """Base diffusion engine storing shared DiT configuration."""
-
-    def __init__(
-        self,
-        dit_config: DiTConfig,
-        model_path: Optional[str] = None,
-        log_stats: bool = False,
-        multiprocess_mode: bool = False,
-    ) -> None:
-        self.dit_config = dit_config
-
-        model_cfg = getattr(dit_config, "model_config", None)
-        config_model_path = None
-        if model_cfg is not None:
-            if isinstance(model_cfg, dict):
-                config_model_path = (
-                    model_cfg.get("model")
-                    or model_cfg.get("model_path")
-                )
-            else:
-                config_model_path = (
-                    getattr(model_cfg, "model", None)
-                    or getattr(model_cfg, "model_path", None)
-                )
-
-        self.model_path = model_path or config_model_path
-        self.model_config = dit_config.model_config
-        self.cache_config = dit_config.dit_cache_config
-
-        self.height = dit_config.height
-        self.width = dit_config.width
-        self.num_inference_steps = dit_config.num_inference_steps
-        self.guidance_scale = dit_config.guidance_scale
-
-        self.log_stats = log_stats
-        self.multiprocess_mode = multiprocess_mode
-
-    def generate(self, *args, **kwargs):  # pragma: no cover - placeholder
-        raise NotImplementedError("This method should be implemented in subclasses.")
-
-
-class DiffusersPipelineEngine(DiffusionEngine):
-    """Adapter that invokes a diffusers pipeline via DiffusionGPUWorker."""
-
-    def __init__(
-        self,
-        dit_config: DiTConfig,
-        model_path: Optional[str] = None,
-        log_stats: bool = False,
-        multiprocess_mode: bool = False,
-    ) -> None:
-        super().__init__(dit_config, model_path, log_stats, multiprocess_mode)
-
-        resolved_model_path = self.model_path
-        if resolved_model_path is None:
-            raise ValueError("Model path must be provided for DiffusersPipelineEngine.")
-
-        device_cfg = getattr(dit_config, "device_config", None)
-        model_cfg = getattr(dit_config, "model_config", None)
-
-        device = None
-        dtype = None
-
-        if device_cfg:
-            if isinstance(device_cfg, dict):
-                device = device_cfg.get("device")
-                dtype = device_cfg.get("dtype")
-            else:
-                device = getattr(device_cfg, "device", None)
-                dtype = getattr(device_cfg, "dtype", None)
-
-        if dtype is None and model_cfg:
-            if isinstance(model_cfg, dict):
-                dtype = model_cfg.get("dtype")
-            else:
-                dtype = getattr(model_cfg, "dtype", None)
-
-        self.worker = DiffusionGPUWorker(
-            resolved_model_path,
-            pipeline_name=dit_config.diffusers_pipeline,
-            device=device,
-            dtype=dtype,
-        )
-
-    def generate(
-        self,
-        prompt: str,
-        *,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_inference_steps: Optional[int] = None,
-        guidance_scale: Optional[float] = None,
-        negative_prompt: Optional[str] = None,
-        seed: Optional[int] = None,
-        image: Optional[Any] = None,
-    ) -> DiffusionRunnerOutput:
-        height = int(height) if height is not None else self.height
-        width = int(width) if width is not None else self.width
-        num_steps = (
-            int(num_inference_steps)
-            if num_inference_steps is not None
-            else self.num_inference_steps
-        )
-        guidance = (
-            float(guidance_scale)
-            if guidance_scale is not None
-            else self.guidance_scale
-        )
-
-        return self.worker.generate(
-            prompt=prompt,
-            height=height,
-            width=width,
-            num_inference_steps=num_steps,
-            guidance_scale=guidance,
-            negative_prompt=negative_prompt,
-            seed=seed,
-            image=image,
-        )

vllm_omni/engine/output_processor.py

@@ -1,21 +1,23 @@
-from typing import Dict, Callable, Optional, Any, Union
+from typing import Any, Callable, Dict, Optional, Union
 
 import torch
 
-from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
-from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestState, RequestOutputCollector
-from vllm.transformers_utils.tokenizer_group import TokenizerGroup
-from vllm.v1.engine import FinishReason
-from vllm.v1.metrics.stats import IterationStats
+from vllm.logger import init_logger
+from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
 from vllm.transformers_utils.tokenizer import AnyTokenizer
-from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
 from vllm.v1.engine.detokenizer import IncrementalDetokenizer
 from vllm.v1.engine.logprobs import LogprobsProcessor
+from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
+from vllm.v1.engine.output_processor import (
+    OutputProcessorOutput,
+    RequestOutputCollector,
+    RequestState,
+)
 from vllm.v1.engine.parallel_sampling import ParentRequest
-from vllm.logger import init_logger
-from vllm.v1.engine import EngineCoreOutput
-
+from vllm.v1.metrics.stats import IterationStats
+from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
 
@@ -56,10 +58,16 @@ class OmniRequestState(RequestState):
                 request=request,
             )
             max_tokens_param = sampling_params.max_tokens
+            top_p = sampling_params.top_p
+            n = sampling_params.n
+            temperature = sampling_params.temperature
         else:
             logprobs_processor = None
             detokenizer = None
             max_tokens_param = None
+            top_p = None
+            n = None
+            temperature = None
             assert request.pooling_params is not None
             output_kind = request.pooling_params.output_kind
 
@@ -67,21 +75,27 @@ class OmniRequestState(RequestState):
             request_id=request.request_id,
             parent_req=parent_req,
             request_index=request_index,
-            lora_name=(request.lora_request.name
-                       if request.lora_request is not None else None),
+            lora_name=(
+                request.lora_request.name if request.lora_request is not None else None
+            ),
             output_kind=output_kind,
             prompt=prompt,
             prompt_token_ids=request.prompt_token_ids,
+            prompt_embeds=request.prompt_embeds,
             logprobs_processor=logprobs_processor,
             detokenizer=detokenizer,
             max_tokens_param=max_tokens_param,
+            top_p=top_p,
+            n=n,
+            temperature=temperature,
             arrival_time=request.arrival_time,
             queue=queue,
             log_stats=log_stats,
         )
 
-    def add_multimodal_tensor(self, tensor: Optional[torch.Tensor],
-                               mm_type: Optional[str]) -> None:
+    def add_multimodal_tensor(
+        self, tensor: Optional[torch.Tensor], mm_type: Optional[str]
+    ) -> None:
         if tensor is None:
             return
         try:
@@ -110,52 +124,49 @@ class OmniRequestState(RequestState):
         finish_reason: Optional[FinishReason],
         stop_reason: Optional[Union[int, str]],
         kv_transfer_params: Optional[dict[str, Any]] = None,
-        num_cached_tokens: Optional[int] = None,
-    ) -> Optional[Any]:
+    ) -> Optional[Union[OmniRequestOutput, PoolingRequestOutput]]:
         finished = finish_reason is not None
         final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
 
         if not finished and final_only:
             return None
 
-        if num_cached_tokens is not None:
-            # Keep num_cached_tokens in RequestOutput for compatibility
-            try:
-                self.num_cached_tokens = num_cached_tokens  # type: ignore[attr-defined]
-            except Exception:
-                pass
-
         request_id = self.request_id
-        output = self._new_completion_output(new_token_ids, finish_reason,
-                                             stop_reason)
+        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)
 
         if self.parent_req is None:
             outputs = [output]
         else:
             request_id, outputs, finished = self.parent_req.get_outputs(
-                request_id, output)
+                request_id, output
+            )
             if not outputs:
                 return None
 
-        return self._new_request_output(request_id, outputs, finished,
-                                        kv_transfer_params)
+        return self._new_request_output(
+            request_id, outputs, finished, kv_transfer_params
+        )
 
     def _new_completion_output(
         self,
         token_ids: list[int],
         finish_reason: Optional[FinishReason],
-        stop_reason: Optional[Union[int, str]]
+        stop_reason: Optional[Union[int, str]],
     ) -> Any:
         # Reuse base text/logprobs logic, then annotate with pooling_result.
-        base_output = super()._new_completion_output(token_ids, finish_reason,
-                                                     stop_reason)
+        base_output = super()._new_completion_output(
+            token_ids, finish_reason, stop_reason
+        )
         try:
             if self.mm_accumulated is not None:
                 tensor = self.mm_accumulated
                 try:
                     tensor = tensor.detach().to("cpu")
                 except Exception:
-                    logger.debug("Failed to move accumulated multimodal tensor to CPU", exc_info=True)
+                    logger.debug(
+                        "Failed to move accumulated multimodal tensor to CPU",
+                        exc_info=True,
+                    )
                 # Attach on the completion output for downstream consumers.
                 if not hasattr(base_output, "multimodal_output"):
                     setattr(base_output, "multimodal_output", {})
@@ -177,17 +188,24 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
       produce the correct RequestOutput/PoolingRequestOutput.
     - Allow custom per-modality handlers via register_handler().
     """
-    def __init__(self, tokenizer: TokenizerGroup, log_stats: bool, engine_core_output_type: Optional[str] = None):
+
+    def __init__(
+        self,
+        tokenizer: AnyTokenizer,
+        log_stats: bool,
+        engine_core_output_type: Optional[str] = None,
+    ):
         super().__init__(tokenizer=tokenizer, log_stats=log_stats)
         self.output_handlers: Dict[str, Callable[[EngineCoreOutput], None]] = {}
         self._reqid_to_mm_type: Dict[str, str] = {}
         self.request_states: dict[str, OmniRequestState] = {}
         self.engine_core_output_type = engine_core_output_type
-        
-    def register_handler(self, modality: str,
-                         handler: Callable[[EngineCoreOutput], None]) -> None:
+
+    def register_handler(
+        self, modality: str, handler: Callable[[EngineCoreOutput], None]
+    ) -> None:
         self.output_handlers[modality.lower()] = handler
-    
+
     def add_request(
         self,
         request: EngineCoreRequest,
@@ -200,21 +218,20 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         if request_id in self.request_states:
             raise ValueError(f"Request id {request_id} already running.")
 
-        tokenizer = None if not self.tokenizer else \
-            self.tokenizer.get_lora_tokenizer(request.lora_request)
-
-        req_state = OmniRequestState.from_new_request(tokenizer=tokenizer,
-                                                  request=request,
-                                                  prompt=prompt,
-                                                  parent_req=parent_req,
-                                                  request_index=request_index,
-                                                  queue=queue,
-                                                  log_stats=self.log_stats)
+        req_state = OmniRequestState.from_new_request(
+            tokenizer=self.tokenizer,
+            request=request,
+            prompt=prompt,
+            parent_req=parent_req,
+            request_index=request_index,
+            queue=queue,
+            log_stats=self.log_stats,
+        )
         self.request_states[request_id] = req_state
         self.lora_states.add_request(req_state)
         if parent_req:
             self.parent_requests[parent_req.request_id] = parent_req
-    
+
     def process_outputs(
         self,
         engine_core_outputs: list[EngineCoreOutput],
@@ -223,7 +240,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
     ) -> OutputProcessorOutput:
         self._reqid_to_mm_type.clear()
         for eco in engine_core_outputs:
-            mm_type = (getattr(eco, "output_type", None) or "").lower()
+            mm_type = (self.engine_core_output_type or "").lower()
             if mm_type:
                 self._reqid_to_mm_type[eco.request_id] = mm_type
             self._route_and_normalize(eco)
@@ -238,23 +255,24 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 continue
 
             # 1) Stats
-            self._update_stats_from_output(req_state, eco,
-                                           engine_core_timestamp,
-                                           iteration_stats)
+            self._update_stats_from_output(
+                req_state, eco, engine_core_timestamp, iteration_stats
+            )
 
             new_token_ids = eco.new_token_ids
             pooling_output = eco.pooling_output
             finish_reason = eco.finish_reason
             stop_reason = eco.stop_reason
             kv_transfer_params = eco.kv_transfer_params
-            num_cached_tokens = eco.num_cached_tokens
+            req_state.num_cached_tokens = eco.num_cached_tokens
             req_state.is_prefilling = False
 
             # 2) Detokenize and logprobs when text path
             assert req_state.detokenizer is not None
             assert req_state.logprobs_processor is not None
             stop_string = req_state.detokenizer.update(
-                new_token_ids, finish_reason == FinishReason.STOP)
+                new_token_ids, finish_reason == FinishReason.STOP
+            )
             if stop_string:
                 finish_reason = FinishReason.STOP
                 stop_reason = stop_string
@@ -262,11 +280,19 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
 
             # 2.5) Accumulate multimodal tensors in RequestState
             try:
-                mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
-                if pooling_output is not None and isinstance(req_state, OmniRequestState):
+                mm_type = (
+                    getattr(eco, "output_type", self.engine_core_output_type) or ""
+                ).lower()
+                if pooling_output is not None and isinstance(
+                    req_state, OmniRequestState
+                ):
                     req_state.add_multimodal_tensor(pooling_output, mm_type)
             except Exception:
-                logger.debug("Failed to accumulate multimodal tensor for request %s", req_id, exc_info=True)
+                logger.debug(
+                    "Failed to accumulate multimodal tensor for request %s",
+                    req_id,
+                    exc_info=True,
+                )
 
             # 3) Create RequestOutput objects, forcing combined mode to keep ids
             pooling_for_make = pooling_output
@@ -274,20 +300,28 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 # Do not consume pooling path now; keep ids and attach mm later
                 pooling_for_make = None
 
-            ro = req_state.make_request_output(new_token_ids, pooling_for_make,
-                                               finish_reason, stop_reason,
-                                               kv_transfer_params,
-                                               num_cached_tokens)
+            ro = req_state.make_request_output(
+                new_token_ids,
+                pooling_for_make,
+                finish_reason,
+                stop_reason,
+                kv_transfer_params,
+            )
             if ro:
                 # Attach accumulated multimodal payload if any
                 try:
-                    if isinstance(req_state, OmniRequestState) and req_state.mm_accumulated is not None:
+                    if (
+                        isinstance(req_state, OmniRequestState)
+                        and req_state.mm_accumulated is not None
+                    ):
                         mm_key = req_state.mm_type or "latents"
                         if not hasattr(ro, "multimodal_output"):
                             setattr(ro, "multimodal_output", {})
                         ro.multimodal_output[mm_key] = req_state.mm_accumulated
                 except Exception:
-                    logger.exception("Error attaching multimodal payload in process_outputs")
+                    logger.exception(
+                        "Error attaching multimodal payload in process_outputs"
+                    )
                 if req_state.queue is not None:
                     req_state.queue.put(ro)
                 else:
@@ -301,8 +335,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                     self.parent_requests.pop(parent_req.request_id, None)
                 if not eco.finished:
                     reqs_to_abort.append(req_id)
-                self._update_stats_from_finished(req_state, finish_reason,
-                                                 iteration_stats)
+                self._update_stats_from_finished(
+                    req_state, finish_reason, iteration_stats
+                )
+                if self.tracer:
+                    self.do_tracing(eco, req_state, iteration_stats)
                 # Cleanup per-request mm state
                 if isinstance(req_state, OmniRequestState):
                     req_state.mm_accumulated = None
@@ -312,10 +349,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             request_outputs=request_outputs,
             reqs_to_abort=reqs_to_abort,
         )
-    
+
     # ---- routing helpers ----
     def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
-        output_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
+        output_type = (
+            getattr(eco, "output_type", self.engine_core_output_type) or ""
+        ).lower()
 
         # Custom handler first (if registered)
         if output_type in self.output_handlers:
@@ -347,7 +386,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         """Ensure image tensors are surfaced via pooling_output for vLLM."""
         if eco.pooling_output is None:
             tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("image", "images", "pixel_values", "pixels"))
+                eco, keys=("image", "images", "pixel_values", "pixels")
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -357,8 +397,17 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         # Preserve text tokens as-is; ensure pooling_output carries image/latents
         if eco.pooling_output is None:
             tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("image", "images", "pixel_values", "pixels",
-                           "latent", "latents", "z"))
+                eco,
+                keys=(
+                    "image",
+                    "images",
+                    "pixel_values",
+                    "pixels",
+                    "latent",
+                    "latents",
+                    "z",
+                ),
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -366,7 +415,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         """Ensure latent tensors are surfaced via pooling_output."""
         if eco.pooling_output is None:
             tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("latent", "latents", "z", "posterior"))
+                eco, keys=("latent", "latents", "z", "posterior")
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -374,8 +424,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         """Ensure audio tensors are surfaced via pooling_output."""
         if eco.pooling_output is None:
             tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("audio", "audios", "wav", "waveform",
-                           "audio_pcm", "pcm"))
+                eco, keys=("audio", "audios", "wav", "waveform", "audio_pcm", "pcm")
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -408,4 +458,4 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         for v in mm.values():
             if isinstance(v, torch.Tensor):
                 return v
-        return None
+        return None

vllm_omni/engine/processor.py

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

vllm_omni/entrypoints/api_server.py

@@ -1,146 +0,0 @@
-"""
-API server for vLLM-omni.
-"""
-
-import asyncio
-from typing import Dict, Any, List
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
-import uvicorn
-
-from .omni_llm import AsyncOmniLLM
-
-
-class GenerateRequest(BaseModel):
-    """Request model for generation."""
-    prompts: List[str]
-    max_tokens: int = 100
-    temperature: float = 0.7
-    stage_args: List[Dict[str, Any]] = None
-
-
-class GenerateResponse(BaseModel):
-    """Response model for generation."""
-    outputs: List[Dict[str, Any]]
-    stage_outputs: List[Dict[str, Any]] = None
-
-
-app = FastAPI(
-    title="vLLM-omni API",
-    description="Multi-modality models inference and serving",
-    version="0.1.0"
-)
-
-# Global omni_llm instance
-omni_llm: AsyncOmniLLM = None
-
-
-@app.on_event("startup")
-async def startup_event():
-    """Initialize the omni_llm instance on startup."""
-    global omni_llm
-    # This will be set by the run_server function
-    pass
-
-
-@app.post("/generate", response_model=GenerateResponse)
-async def generate(request: GenerateRequest):
-    """Generate text or multimodal content."""
-    try:
-        if omni_llm is None:
-            raise HTTPException(status_code=500, detail="OmniLLM not initialized")
-        
-        # Prepare stage arguments
-        if request.stage_args is None:
-            # Create default stage arguments - one per stage config
-            # For now, we'll process all prompts in the first stage
-            stage_args = [{
-                "prompt": " ".join(request.prompts) if request.prompts else "",
-                "max_tokens": request.max_tokens,
-                "temperature": request.temperature
-            }]
-        else:
-            stage_args = request.stage_args
-        
-        # Generate using omni_llm
-        try:
-            outputs = await omni_llm.generate_async(stage_args)
-        except Exception as e:
-            # Fallback to synchronous generation
-            outputs = omni_llm.generate(stage_args)
-        
-        # Convert outputs to response format
-        response_outputs = []
-        for output in outputs:
-            if hasattr(output, 'outputs') and output.outputs:
-                for out in output.outputs:
-                    response_outputs.append({
-                        "text": getattr(out, 'text', ''),
-                        "finished": getattr(out, 'finish_reason', 'length') != 'length',
-                        "tokens": getattr(out, 'token_ids', [])
-                    })
-            else:
-                response_outputs.append({
-                    "text": "",
-                    "finished": True,
-                    "tokens": []
-                })
-        
-        return GenerateResponse(
-            outputs=response_outputs,
-            stage_outputs=[{"stage": i, "output": "processed"} for i in range(len(stage_args))]
-        )
-    
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-
-@app.get("/health")
-async def health_check():
-    """Health check endpoint."""
-    return {"status": "healthy", "service": "vllm-omni"}
-
-
-@app.get("/info")
-async def get_info():
-    """Get information about the service."""
-    if omni_llm is None:
-        return {"error": "OmniLLM not initialized"}
-    
-    return {
-        "service": "vllm-omni",
-        "version": "0.1.0",
-        "num_stages": omni_llm.get_num_stages() if hasattr(omni_llm, 'get_num_stages') else 0,
-        "stage_configs": [
-            {
-                "stage_id": config.stage_id,
-                "engine_type": config.engine_type,
-                "model_path": config.model_path,
-                "input_modalities": config.input_modalities,
-                "output_modalities": config.output_modalities
-            }
-            for config in omni_llm.stage_configs
-        ]
-    }
-
-
-async def run_server(omni_llm_instance: AsyncOmniLLM, host: str = "0.0.0.0", port: int = 8000):
-    """Run the API server."""
-    global omni_llm
-    omni_llm = omni_llm_instance
-    
-    config = uvicorn.Config(
-        app=app,
-        host=host,
-        port=port,
-        log_level="info"
-    )
-    
-    server = uvicorn.Server(config)
-    await server.serve()
-
-
-if __name__ == "__main__":
-    # This is for testing purposes
-    asyncio.run(run_server(None))

vllm_omni/entrypoints/cli/main.py

@@ -3,8 +3,7 @@ CLI entry point for vLLM-omni that intercepts vLLM commands.
 """
 
 import sys
-import argparse
-from typing import List, Optional
+
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
 
 
@@ -19,6 +18,7 @@ def main():
     else:
         # Forward to original vLLM CLI
         from vllm.entrypoints.cli.main import main as vllm_main
+
         vllm_main()
 
 

vllm_omni/entrypoints/cli/serve.py

@@ -3,204 +3,91 @@ Omni serve command for vLLM-omni.
 """
 
 import argparse
-import asyncio
-from typing import List, Optional
-
-from ..omni_llm import AsyncOmniLLM
-from vllm_omni.config import (
-    DiTConfig,
-    create_ar_stage_config,
-    create_dit_stage_config,
-)
+from typing import List
 
 
 class OmniServeCommand:
     """Command handler for vLLM-omni serve command."""
-    
+
     def __init__(self):
         self.parser = self._create_parser()
-    
+
     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"
         )
-        
+
         # 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)"
+            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"
+            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"
+            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("--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"
+            help="Number of DiT inference steps",
         )
-        
+
         parser.add_argument(
             "--dit-guidance-scale",
             type=float,
             default=7.5,
-            help="DiT guidance scale"
+            help="DiT guidance scale",
         )
-        
+
         parser.add_argument(
             "--use-diffusers",
             action="store_true",
-            help="Use diffusers pipeline for DiT stage"
+            help="Use diffusers pipeline for DiT stage",
         )
-        
+
         # Other arguments
         parser.add_argument(
-            "--log-stats",
-            action="store_true",
-            help="Enable logging statistics"
+            "--log-stats", action="store_true", help="Enable logging statistics"
         )
-        
+
         return parser
-    
+
     def run(self, args: List[str]) -> None:
         """Run the omni serve command."""
-        parsed_args = self.parser.parse_args(args)
-        
-        # Create stage configurations
-        stage_configs = self._create_stage_configs(parsed_args)
-        
-        # Create AsyncOmniLLM instance
-        omni_llm = AsyncOmniLLM(
-            stage_configs=stage_configs,
-            log_stats=parsed_args.log_stats
-        )
-        
-        # Start the server
-        asyncio.run(self._start_server(omni_llm, parsed_args))
-    
+        pass
+
     def _create_stage_configs(self, args) -> List:
         """Create stage configurations based on arguments."""
-        stage_configs = []
-        stage_id = 0
-        
-        # Add AR stage - use main model if ar_stage not specified
-        ar_model = args.ar_stage if args.ar_stage else args.model
-        ar_config = create_ar_stage_config(
-            stage_id=stage_id,
-            model_path=ar_model,
-            input_modalities=["text"],
-            output_modalities=["text"]
-        )
-        stage_configs.append(ar_config)
-        stage_id += 1
-        
-        # Add DiT stage if specified
-        if args.dit_stage:
-            dit_model = args.dit_stage
-        elif args.use_diffusers:
-            # Use a default DiT model if diffusers is enabled
-            dit_model = "stabilityai/stable-diffusion-2-1"
-        else:
-            dit_model = None
-            
-        if dit_model:
-            dit_config = DiTConfig(
-                num_inference_steps=args.dit_steps,
-                guidance_scale=args.dit_guidance_scale,
-                use_diffusers=args.use_diffusers,
-                diffusers_pipeline="auto" if args.use_diffusers else None,
-            )
-            
-            dit_stage_config = create_dit_stage_config(
-                stage_id=stage_id,
-                model_path=dit_model,
-                input_modalities=["text"],
-                output_modalities=["image"],
-                dit_config=dit_config
-            )
-            stage_configs.append(dit_stage_config)
-            stage_id += 1
-        
-        # If no specific stages are specified, use the main model
-        if not stage_configs:
-            # Try to detect if it's a multimodal model
-            if "omni" in args.model.lower() or "multimodal" in args.model.lower():
-                # Assume it's a multimodal model that can handle both AR and DiT
-                ar_config = create_ar_stage_config(
-                    stage_id=0,
-                    model_path=args.model,
-                    input_modalities=["text"],
-                    output_modalities=["text"]
-                )
-                stage_configs.append(ar_config)
-                
-                dit_config = DiTConfig(
-                    num_inference_steps=args.dit_steps,
-                    guidance_scale=args.dit_guidance_scale,
-                    use_diffusers=args.use_diffusers,
-                    diffusers_pipeline="auto" if args.use_diffusers else None,
-                )
-                
-                dit_stage_config = create_dit_stage_config(
-                    stage_id=1,
-                    model_path=args.model,
-                    input_modalities=["text"],
-                    output_modalities=["image"],
-                    dit_config=dit_config
-                )
-                stage_configs.append(dit_stage_config)
-            else:
-                # Default to AR stage
-                ar_config = create_ar_stage_config(
-                    stage_id=0,
-                    model_path=args.model,
-                    input_modalities=["text"],
-                    output_modalities=["text"]
-                )
-                stage_configs.append(ar_config)
-        
-        return stage_configs
-    
-    async def _start_server(self, omni_llm: AsyncOmniLLM, args) -> None:
+        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
-            )
+
+            await run_server(omni_llm_instance=omni_llm, host=args.host, port=args.port)
         except KeyboardInterrupt:
             print("\nShutting down server...")
         except Exception as e:

vllm_omni/entrypoints/log_utils.py

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

vllm_omni/entrypoints/omni_llm.py

@@ -1,88 +1,398 @@
-from typing import Union, Sequence, Optional, Any
+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
+
 import cloudpickle
 from pydantic import ValidationError
 
-from vllm.inputs import PromptType
-from vllm.sampling_params import SamplingParams
-from vllm.entrypoints.llm import LLM
-
-from vllm.v1.engine.llm_engine import LLMEngine
+# External library imports (vLLM)
+from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field
 from vllm.engine.arg_utils import HfOverrides
-from vllm.usage.usage_lib import UsageContext
-from vllm.config import CompilationConfig, is_init_field
-from vllm.utils import Counter
+from vllm.entrypoints.llm import LLM
+from vllm.inputs import PromptType
 from vllm.logger import init_logger
-import vllm.envs as envs
+from vllm.plugins.io_processors import get_io_processor
+from vllm.sampling_params import SamplingParams
+from vllm.usage.usage_lib import UsageContext
+from vllm.utils import Counter
+from vllm.v1.engine.llm_engine import LLMEngine
 
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model
-from vllm_omni.entrypoints.omni_stage import OmniStage
+# Internal imports (our code)
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
+from vllm_omni.entrypoints.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,
+    configure_orchestrator_logger,
+    init_stats_paths,
+    OrchestratorMetrics,
+)
+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.outputs import OmniRequestOutput
 
-
 logger = init_logger(__name__)
 
 
 class OmniLLM:
-    def __init__(self, model: str, stage_configs = None, log_stats: bool = False, **kwargs):
+    """Multi-process pipelined OmniLLM.
+
+    - Per-stage process with copy-based IPC (Queues)
+    - window=-1 across stages (downstream starts when upstream finishes)
+    - max_inflight=1 per stage (serial within a stage), but pipeline across stages
+    """
+
+    def __init__(self, model: str,
+                 stage_configs=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:
-            self.initialize_stage_configs(model)
+            self.stage_configs = load_stage_configs_from_model(model)
         else:
             self.stage_configs = stage_configs
+
+
+        # Optional file handler for orchestrator
+        self._log_file = log_file
+        if self._log_file:
+            remove_old_logs(self._log_file, len(self.stage_list))
+            configure_orchestrator_logger(logger, self._log_file)
+
+        self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
+        self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
+
+    def _initialize_stages(self, model: str, init_sleep_seconds: int, shm_threshold_bytes: int, init_timeout: int) -> None:
+        self.stage_list: list[OmniStage] = []
+        # Build OmniStage instances in parallel, preserve original order
+        def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
+            idx, cfg = idx_cfg
+            return idx, OmniStage(cfg)
+
+        with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor:
+            futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)]
+            results: list[tuple[int, OmniStage]] = []
+            for fut in as_completed(futures):
+                results.append(fut.result())
+        results.sort(key=lambda x: x[0])
+        self.stage_list = [st for _, st in results]
+        logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
+
+        self._ctx = mp.get_context("spawn")
+        self._stage_in_queues: list[mp.Queue] = []
+        self._stage_out_queues: list[mp.Queue] = []
+        self._init_sleep_seconds = max(0, int(init_sleep_seconds))
+        self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
+        self._start_stage_processes(model)
+        # Wait for all stages to report readiness before seeding
+        self._stages_ready: set[int] = set()
+        self._wait_for_stages_ready(timeout=init_timeout)
         
-        self.stage_list = []
-        self.initialize_stages(model)
-        
-    def initialize_stage_configs(self, model: str):
-        self.stage_configs = load_stage_configs_from_model(model)
-    
-    def initialize_stages(self, model: str):
-        for stage_config in self.stage_configs:
-            stage = OmniStage(stage_config)
-            stage_llm = OmniStageLLM(model=model, **stage_config.engine_args)
-            stage.set_engine(stage_llm)
-            self.stage_list.append(stage)
-    
+
+    def _start_stage_processes(self, model: str) -> None:
+        for stage_id, stage in enumerate(self.stage_list):
+            # Use unbounded queues to avoid deadlock when seeding many requests
+            in_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            self._stage_in_queues.append(in_q)
+            self._stage_out_queues.append(out_q)
+
+            # Attach queues and start Stage-owned worker process
+            stage.attach_queues(in_q, out_q)
+            stage.init_stage_worker(
+                model,
+                log_file=self._log_file,
+                shm_threshold_bytes=self._shm_threshold_bytes,
+                ctx=self._ctx,
+                batch_timeout=self.batch_timeout,
+            )
+            logger.debug("[Orchestrator] Stage-%s process started", stage_id)
+            time.sleep(self._init_sleep_seconds)
+
+    def close(self) -> None:
+        for q in self._stage_in_queues:
+            try:
+                q.put_nowait(None)
+            except Exception as e:
+                logger.warning("[Orchestrator] Failed to send shutdown signal to stage input queue: %s", e)
+        for stage in self.stage_list:
+            try:
+                stage.stop_stage_worker()
+            except Exception as e:
+                logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
+
+    def __del__(self) -> None:  # best-effort
+        try:
+            self.close()
+        except Exception as e:
+            logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
+
     def generate(
         self,
         prompts: Union[PromptType, Sequence[PromptType]],
-        sampling_params_list: Optional[Union[SamplingParams,
-                                        Sequence[SamplingParams]]] = None,
+        sampling_params_list: Optional[
+            Union[SamplingParams, Sequence[SamplingParams]]
+        ] = None,
     ) -> list[OmniRequestOutput]:
-        """Generate text outputs for the given prompts."""
-        final_outputs: list[OmniRequestOutput] = []
+        logger.debug("[Orchestrator] generate() called")
+        if sampling_params_list is None:
+            raise ValueError("sampling_params_list is required for pipelined generation")
         if len(sampling_params_list) != len(self.stage_list):
-            raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
-        for stage_id, stage in enumerate(self.stage_list):
-            if stage_id > 0:
-                engine_inputs = stage.process_engine_inputs(self.stage_list, prompts)
-            else:
-                engine_inputs = prompts
-            engine_outputs = self._run_generation(stage, sampling_params_list[stage_id], engine_inputs)
-            stage.set_engine_outputs(engine_outputs)
-            if hasattr(stage, 'final_output') and stage.final_output:
-                final_outputs.append(OmniRequestOutput(
-                    stage_id=stage_id, 
-                    final_output_type=stage.final_output_type, 
-                    request_output=engine_outputs))
+            raise ValueError(
+                f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}"
+            )
+
+        # Normalize prompts to a list for per-request iteration
+        if not isinstance(prompts, (list, tuple)):
+            request_prompts: list[PromptType] = [prompts]
+        else:
+            request_prompts = list(prompts)
+
+        final_outputs: list[OmniRequestOutput] = []
+
+        # Orchestrator keeps stage objects for input derivation
+        num_stages = len(self.stage_list)
+
+        # Map from request_id to original prompt
+        request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
+
+        # Track per-request start time for end-to-end timing
+        _req_start_ts: dict[int, float] = {}
+        _wall_start_ts: float = time.time()
+        _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 %d requests into stage-0", len(request_prompts))
+        # Mark first input time for stage-0
+        metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
+
+        for req_id, prompt in request_id_to_prompt.items():
+            sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
+            task = {
+                "request_id": req_id,
+                "engine_inputs": prompt,
+                "sampling_params": sp0,
+            }
+            self.stage_list[0].submit(task)
+            _req_start_ts[req_id] = time.time()
+            logger.debug("[Orchestrator] Enqueued request %s to stage-0", req_id)
+
+        # For each stage, forward results to next stage; collect finals at the end
+        # We pipeline by continually polling output queues in stage order
+        remaining_by_stage: list[int] = [len(request_prompts)] + [0] * (num_stages - 1)
+        completed_requests = 0
+        total_requests = len(request_prompts)
+
+        logger.debug("[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d", total_requests, num_stages)
+        while completed_requests < total_requests:
+            made_progress = False
+            for stage_id, stage in enumerate(self.stage_list):
+                result = stage.try_collect()
+                if result is None:
+                    continue
+
+                made_progress = True
+                req_id = result.get("request_id")
+                if "error" in result:
+                    logger.error("Stage %s error on request %s: %s", stage_id, req_id, result["error"])
+                    continue
+                
+                if result.get("type") == "stage_ready":
+                    #Only happens when stage is initialized slower than expected, so we wait for a short time and try again
+                    time.sleep(0.05)
+                    continue
+
+                engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
+                # Mark last output time for this stage whenever we receive outputs
+                metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
+                try:
+                    _m = result.get("metrics")
+                    if _m is not None:
+                        metrics.on_stage_metrics(stage_id, req_id, _m)
+                except Exception as e:
+                    logger.exception("[Orchestrator] Failed to process metrics for stage %s, req %s: %s", stage_id, req_id, e)
+                logger.debug("[Orchestrator] Stage-%s completed request %s; forwarding or finalizing", stage_id, req_id)
+                stage.set_engine_outputs(engine_outputs)
+
+                if getattr(stage, "final_output", False):
+                    final_outputs.append(
+                        OmniRequestOutput(
+                            stage_id=stage_id,
+                            final_output_type=stage.final_output_type,  # type: ignore[attr-defined]
+                            request_output=engine_outputs,
+                        )
+                    )
+                    logger.debug("[Orchestrator] Request %s finalized at stage-%s", req_id, stage_id)
+
+                    # End-to-end timing and time-per-token for final output (only once per request at the designated final stage)
+                    try:
+                        rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
+                        if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
+                            metrics.on_finalize_request(stage_id, req_id, engine_outputs, _req_start_ts.get(req_id, _wall_start_ts))
+                    except Exception as e:
+                        logger.exception("[Orchestrator] Finalize request handling error for req %s at stage %s: %s", req_id, stage_id, e)
+
+                next_stage_id = stage_id + 1
+                if next_stage_id < num_stages:
+                    next_stage: OmniStage = self.stage_list[next_stage_id]
+                    next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
+                    sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
+                    try:
+                        # Measure transfer size and time (encode + enqueue)
+                        size_bytes = 0
+                        try:
+                            size_bytes = len(_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)
+                    remaining_by_stage[next_stage_id] += 1
+                else:
+                    completed_requests += 1
+                    logger.debug("[Orchestrator] Request %s fully completed (%d/%d)", req_id, completed_requests, total_requests)
+
+            if not made_progress:
+                time.sleep(0.005)
+        logger.debug("[Orchestrator] All requests completed")
+
+        # Summarize and print stats
+        try:
+            summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
+            logger.info("[Summary] %s", summary)
+        except Exception as e:
+            logger.exception("[Orchestrator] Failed to build/log summary: %s", e)
+
         return final_outputs
-    
-    def _run_generation(self, stage: OmniStage, sampling_params: SamplingParams, prompts: Union[PromptType, Sequence[PromptType]]):
-        engine_outputs = []
-        for ro in stage.engine.generate(prompts, sampling_params):
-            engine_outputs.append(ro)
-        return engine_outputs
+
+    def _wait_for_stages_ready(self, timeout: int = 120) -> None:
+        deadline = time.time() + max(0, int(timeout))
+        num_stages = len(self.stage_list)
+        while len(self._stages_ready) < num_stages and time.time() < deadline:
+            progressed = False
+            for stage_id, stage in enumerate(self.stage_list):
+                if stage_id in self._stages_ready:
+                    continue
+                result = stage.try_collect()
+                if result is None:
+                    continue
+                progressed = True
+                if result.get("type") == "stage_ready":
+                    self._stages_ready.add(stage_id)
+                    logger.debug("[Orchestrator] Stage-%s reported ready", stage_id)
+                else:
+                    # No user data should arrive before seeding; ignore other messages
+                    pass
+            if not progressed:
+                time.sleep(0.01)
+        if len(self._stages_ready) < num_stages:
+            not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
+            logger.warning(
+                "[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s",
+                len(self._stages_ready), num_stages, not_ready,
+            )
+            # Provide actionable suggestions before shutdown
+            try:
+                suggestions = [
+                    "Verify GPU/device assignment in config (runtime.devices) is correct.",
+                    "Check GPU/host memory availability; reduce model or batch size if needed.",
+                    "Check model weights path and network reachability (if loading remotely).",
+                    "Increase initialization wait time (init_sleep_seconds or call-site timeout).",
+                ]
+                if getattr(self, "_log_file", None):
+                    suggestions.append(
+                        f"Inspect per-stage log files for details: {self._log_file}.stage<id>.log"
+                    )
+                logger.error(
+                    "[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s",
+                    "\n- ".join(suggestions),
+                )
+            except Exception:
+                # Best-effort logging of suggestions
+                logger.error(
+                    "[Orchestrator] Stage initialization failed and an error occurred while logging suggestions",
+                )
+
+            # Attempt graceful shutdown of all stages before exiting
+            try:
+                self.close()
+            except Exception:
+                pass
+
+            # Terminate the current process with non-zero exit code
+            try:
+                sys.exit(1)
+            except SystemExit:
+                raise
+            except Exception:
+                os._exit(1)
 
 
-class OmniStageLLM(LLM): 
-    def __init__(self, 
-                model: str, 
-                compilation_config: Optional[Union[int, dict[str, Any],
-                                           CompilationConfig]] = None,
-                hf_overrides: Optional[HfOverrides] = None,
-                **kwargs):
+class OmniStageLLM(LLM):
+    def __init__(
+        self,
+        model: str,
+        compilation_config: Optional[
+            Union[int, dict[str, Any], CompilationConfig]
+        ] = None,
+        hf_overrides: Optional[HfOverrides] = None,
+        structured_outputs_config: Optional[
+            Union[dict[str, Any], StructuredOutputsConfig]
+        ] = None,
+        **kwargs,
+    ):
         """LLM constructor."""
         if "disable_log_stats" not in kwargs:
             kwargs["disable_log_stats"] = True
@@ -95,21 +405,23 @@ class OmniStageLLM(LLM):
                 kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)
 
         if "kv_transfer_config" in kwargs and isinstance(
-                kwargs["kv_transfer_config"], dict):
-            from vllm.config import KVTransferConfig
+            kwargs["kv_transfer_config"], dict
+        ):
+            from vllm.config.kv_transfer import KVTransferConfig
+
             raw_config_dict = kwargs["kv_transfer_config"]
             try:
-                kwargs["kv_transfer_config"] = KVTransferConfig(
-                    **raw_config_dict)
+                kwargs["kv_transfer_config"] = KVTransferConfig(**raw_config_dict)
             except ValidationError as e:
                 logger.error(
                     "Failed to convert 'kv_transfer_config' dict to "
                     "KVTransferConfig object. Dict: %s. Error: %s",
-                    raw_config_dict, e)
+                    raw_config_dict,
+                    e,
+                )
                 # Consider re-raising a more specific vLLM error or ValueError
                 # to provide better context to the user.
-                raise ValueError(
-                    f"Invalid 'kv_transfer_config' provided: {e}") from e
+                raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e
 
         if hf_overrides is None:
             hf_overrides = {}
@@ -117,41 +429,68 @@ class OmniStageLLM(LLM):
         if compilation_config is not None:
             if isinstance(compilation_config, int):
                 compilation_config_instance = CompilationConfig(
-                    level=compilation_config)
+                    level=compilation_config
+                )
             elif isinstance(compilation_config, dict):
-                predicate = lambda x: is_init_field(CompilationConfig, x[0])
                 compilation_config_instance = CompilationConfig(
-                    **dict(filter(predicate, compilation_config.items())))
+                    **{
+                        k: v
+                        for k, v in compilation_config.items()
+                        if is_init_field(CompilationConfig, k)
+                    }
+                )
             else:
                 compilation_config_instance = compilation_config
         else:
             compilation_config_instance = CompilationConfig()
 
+        if structured_outputs_config is not None:
+            if isinstance(structured_outputs_config, dict):
+                structured_outputs_instance = StructuredOutputsConfig(
+                    **{
+                        k: v
+                        for k, v in structured_outputs_config.items()
+                        if is_init_field(StructuredOutputsConfig, k)
+                    }
+                )
+            else:
+                structured_outputs_instance = structured_outputs_config
+        else:
+            structured_outputs_instance = StructuredOutputsConfig()
+
         engine_args = OmniEngineArgs(
             model=model,
             hf_overrides=hf_overrides,
             compilation_config=compilation_config_instance,
+            structured_outputs_config=structured_outputs_instance,
             **kwargs,
         )
 
         # Create the Engine (autoselects V0 vs V1)
         self.llm_engine = LLMEngine.from_engine_args(
-            engine_args=engine_args, usage_context=UsageContext.LLM_CLASS)
-        self.llm_engine.output_processor = MultimodalOutputProcessor(tokenizer=self.llm_engine.tokenizer, 
-                                                                    log_stats=self.llm_engine.log_stats)
-        self.llm_engine.processor = OmniProcessor(vllm_config=self.llm_engine.vllm_config,
-                                                  tokenizer=self.llm_engine.tokenizer)
+            engine_args=engine_args, usage_context=UsageContext.LLM_CLASS
+        )
+        self.llm_engine.output_processor = MultimodalOutputProcessor(
+            tokenizer=self.llm_engine.tokenizer,
+            log_stats=self.llm_engine.log_stats,
+            engine_core_output_type=engine_args.engine_output_type,
+        )
+        self.llm_engine.processor = OmniProcessor(
+            vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer
+        )
         self.engine_class = type(self.llm_engine)
 
         self.request_counter = Counter()
         self.default_sampling_params: Union[dict[str, Any], None] = None
 
-        if envs.VLLM_USE_V1:
-            supported_tasks = self.llm_engine \
-                .get_supported_tasks()  # type: ignore
-        else:
-            supported_tasks = self.llm_engine.model_config.supported_tasks
+        supported_tasks = self.llm_engine.get_supported_tasks()  # type: ignore
 
         logger.info("Supported_tasks: %s", supported_tasks)
 
-        self.supported_tasks = supported_tasks
+        self.supported_tasks = supported_tasks
+
+        # Load the Input/Output processor plugin if any
+        io_processor_plugin = self.llm_engine.model_config.io_processor_plugin
+        self.io_processor = get_io_processor(
+            self.llm_engine.vllm_config, io_processor_plugin
+        )

vllm_omni/entrypoints/omni_stage.py

@@ -1,14 +1,27 @@
 """
 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.
 """
 
 import importlib
-from typing import List, Union
-from vllm.v1.engine.llm_engine import LLMEngine
-from vllm.v1.engine.async_llm import AsyncLLM
+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.v1.engine import EngineCoreOutput
-
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.engine.llm_engine import LLMEngine
 from vllm_omni.inputs.data import OmniTokensPrompt
 
 
@@ -20,30 +33,40 @@ class OmniStage:
         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'):
+        if hasattr(stage_config, "engine_input_source"):
             self.engine_input_source = stage_config.engine_input_source
         else:
             self.engine_input_source = []
         self.engine_output_type = stage_config.engine_args.engine_output_type
         self.engine_outputs = None
-        if hasattr(stage_config, 'custom_process_input_func'):
+        if hasattr(stage_config, "custom_process_input_func"):
             # Import the module specified in the config (already a full module path)
-            module_path, func_name = stage_config.custom_process_input_func.rsplit('.', 1)
+            module_path, func_name = stage_config.custom_process_input_func.rsplit(
+                ".", 1
+            )
             module = importlib.import_module(module_path)
             self.custom_process_input_func = getattr(module, func_name)
         else:
             self.custom_process_input_func = None
 
-        if hasattr(stage_config, 'final_output'):
+        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'):
+        if hasattr(stage_config, "final_output_type"):
             self.final_output_type = stage_config.final_output_type
         else:
             self.final_output_type = None
 
+        # Runtime orchestration state (added)
+        self._in_q: Optional[mp.Queue] = None
+        self._out_q: Optional[mp.Queue] = None
+        self._proc: Optional[mp.Process] = None
+        self._log_file: Optional[str] = None
+        self._shm_threshold_bytes: int = 65536
+        self._logger = logging.getLogger(__name__)
+
     def set_engine(self, engine: LLMEngine) -> None:
         """Initialize the engine for the stage."""
         self.engine = engine
@@ -51,12 +74,75 @@ class OmniStage:
     def set_async_engine(self, async_engine: AsyncLLM) -> None:
         """Initialize the async engine for the stage."""
         self.async_engine = async_engine
-    
+
     def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None:
         """Set the engine output for the stage."""
         self.engine_outputs = engine_outputs
-    
-    def process_engine_inputs(self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None) -> List[Union[OmniTokensPrompt, TextPrompt]]:
+
+    # ----------------- New Orchestration APIs -----------------
+    def attach_queues(self, in_q: mp.Queue, out_q: mp.Queue) -> None:
+        self._in_q = in_q
+        self._out_q = out_q
+
+    def init_stage_worker(
+        self,
+        model: str,
+        *,
+        log_file: Optional[str] = None,
+        shm_threshold_bytes: int = 65536,
+        ctx: Optional[mp.context.BaseContext] = None,
+        batch_timeout: int = 10,
+    ) -> None:
+        assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process"
+        self._log_file = log_file
+        self._shm_threshold_bytes = shm_threshold_bytes
+        ctx = ctx or mp.get_context("spawn")
+        # Prepare lightweight dict config for worker
+        engine_args = _to_dict(self.engine_args)
+        runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {}))
+        stage_payload: Dict[str, Any] = {
+            "stage_id": self.stage_id,
+            "engine_args": engine_args,
+            "runtime": runtime_cfg,
+            "shm_threshold_bytes": self._shm_threshold_bytes,
+        }
+        self._proc = ctx.Process(
+            target=_stage_worker,
+            args=(model, stage_payload, self._in_q, self._out_q, self._log_file, batch_timeout),
+        )
+        self._proc.start()
+
+    def stop_stage_worker(self) -> None:
+        if self._in_q is not None:
+            try:
+                self._in_q.put_nowait(None)
+            except Exception as e:
+                self._logger.warning("[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e)
+        if self._proc is not None:
+            try:
+                self._proc.join(timeout=5)
+            except Exception as e:
+                self._logger.debug("[Stage-%s] join() failed: %s", self.stage_id, e, exc_info=True)
+            if self._proc.is_alive():
+                try:
+                    self._proc.terminate()
+                except Exception as e:
+                    self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
+
+    def submit(self, payload: Dict[str, Any]) -> None:
+        assert self._in_q is not None
+        self._in_q.put(payload)
+
+    def try_collect(self) -> Optional[Dict[str, Any]]:
+        assert self._out_q is not None
+        try:
+            return self._out_q.get_nowait()
+        except Exception:
+            return None
+
+    def process_engine_inputs(
+        self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None
+    ) -> List[Union[OmniTokensPrompt, TextPrompt]]:
         """Process the engine input for the stage."""
         if self.custom_process_input_func is None:
             engine_inputs = []
@@ -64,17 +150,246 @@ 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
-            multi_modal_data = {source_output.request_id: 
-            prompt.get('multi_modal_data', None) for source_output, prompt in zip(source_outputs, prompt)}
+            multi_modal_data = {
+                source_output.request_id: prompt.get("multi_modal_data", None)
+                for source_output, prompt in zip(source_outputs, prompt)
+            }
 
             for source_output in source_outputs:
                 engine_input = OmniTokensPrompt(
-                    prompt_token_ids = source_output.outputs[0].token_ids,
-                    multi_modal_data=multi_modal_data[source_output.request_id] if multi_modal_data else None,
+                    prompt_token_ids=source_output.outputs[0].token_ids,
+                    multi_modal_data=(
+                        multi_modal_data[source_output.request_id]
+                        if multi_modal_data
+                        else None
+                    ),
                 )
                 engine_inputs.append(engine_input)
             return engine_inputs
 
         else:
             engine_input_source = self.engine_input_source
-            return self.custom_process_input_func(stage_list, engine_input_source, prompt)
+            return self.custom_process_input_func(
+                stage_list, engine_input_source, prompt
+            )
+
+
+def _stage_worker(
+    model: str,
+    stage_payload: Dict[str, Any],
+    in_q: mp.Queue,
+    out_q: mp.Queue,
+    log_file: Optional[str] = None,
+    batch_timeout: int = 10,
+) -> None:
+    """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
+    import logging as _logging
+    from vllm_omni.entrypoints.omni_llm import OmniStageLLM  # noqa: WPS433
+    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,
+    )
+    import queue as _queue
+    import time as _time
+    # no inline JSONL/serialization imports; logging handled by utilities
+
+    stage_id = stage_payload["stage_id"]
+    engine_args = stage_payload.get("engine_args", {})
+    runtime_cfg = stage_payload.get("runtime", {})
+    shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
+
+    # Per-stage file logger (optional)
+    try:
+        if log_file:
+            stage_log = _logging.getLogger(__name__)
+            stage_log.setLevel(_logging.DEBUG)
+            fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
+            fh.setLevel(_logging.DEBUG)
+            fh.setFormatter(_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s"))
+            class _StageFilter(_logging.Filter):
+                def filter(self, record: _logging.LogRecord) -> bool:
+                    setattr(record, "stage", stage_id)
+                    return True
+            fh.addFilter(_StageFilter())
+            stage_log.addHandler(fh)
+    except Exception:
+        pass
+
+    # Stage stats JSONL file
+    _stats_file = f"{log_file}.stage{stage_id}.stats.jsonl" if log_file else None
+
+    # Aggregates for running average
+    _agg_total_tokens = 0
+    _agg_total_gen_time_ms = 0.0
+    # Monotonic batch id per stage process for orchestrator dedup on time aggregation
+    _batch_seq = 0
+
+    # Device mapping
+    try:
+        set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
+    except Exception as e:
+        _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
+
+    # Init LLM
+    _logging.getLogger(__name__).debug("[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys()))
+    stage_engine = OmniStageLLM(model=model, **engine_args)
+    _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
+    # Signal readiness to orchestrator
+    try:
+        out_q.put({"type": "stage_ready", "stage_id": stage_id})
+    except Exception:
+        pass
+
+    # Batch processing loop
+    while True:
+        task = in_q.get()
+        _recv_dequeue_ts = _time.time()
+        if task is None:
+            _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id)
+            break
+
+        max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
+        batch_tasks: List[Dict[str, Any]] = [task]
+        if max_batch_size > 1:
+            while len(batch_tasks) < max_batch_size:
+                if not in_q.empty():
+                    extra = in_q.get(timeout=batch_timeout)
+                    if extra is None:
+                        in_q.put(None)
+                        break
+                    batch_tasks.append(extra)
+                else:
+                    break
+
+        batch_request_ids: List[Any] = []
+        batch_engine_inputs: List[Any] = []
+        _rx_bytes_by_rid: Dict[Any, int] = {}
+        _rx_decode_ms_by_rid: Dict[Any, float] = {}
+        _in_flight_ms_by_rid: Dict[Any, float] = {}
+        for t in batch_tasks:
+            rid = t["request_id"]
+            try:
+                sent_ts = float(t.get("sent_ts", None)) if isinstance(t, dict) else None
+                if sent_ts is not None:
+                    _in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0
+                else:
+                    _in_flight_ms_by_rid[rid] = 0.0
+            except Exception:
+                _in_flight_ms_by_rid[rid] = 0.0
+            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(
+                t, obj_key="engine_inputs", shm_key="engine_inputs_shm"
+            )
+            _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
+            _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
+            batch_request_ids.append(rid)
+            if isinstance(ein, list):
+                batch_engine_inputs.extend(ein)
+            elif isinstance(ein, dict):
+                batch_engine_inputs.append(ein)
+            else:
+                _logging.getLogger(__name__).exception("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein))
+        sampling_params = batch_tasks[0]["sampling_params"]
+        _logging.getLogger(__name__).debug("[Stage-%s] Received batch size=%d, request_ids=%s", stage_id, len(batch_tasks), batch_request_ids)
+        print("--------------------------------", flush=True)
+        print(f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}", flush=True)
+        print("--------------------------------", flush=True)
+        try:
+            _batch_seq += 1
+            gen_outputs: List[Any] = []
+            _gen_t0 = _time.time()
+            for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False):
+                gen_outputs.append(ro)
+            _gen_t1 = _time.time()
+            _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
+
+            # Group outputs per request id with fallback
+            req_to_outputs: Dict[Any, List[Any]] = {rid: [] for rid in batch_request_ids}
+            unmapped: List[Any] = []
+            for ro in gen_outputs:
+                rid = getattr(ro, "request_id", None)
+                if rid in req_to_outputs:
+                    req_to_outputs[rid].append(ro)
+                else:
+                    unmapped.append(ro)
+            if unmapped:
+                idx = 0
+                for ro in unmapped:
+                    target_rid = batch_request_ids[idx % len(batch_request_ids)]
+                    req_to_outputs[target_rid].append(ro)
+                    idx += 1
+
+            # Per-request stats logging and aggregates
+            for rid in batch_request_ids:
+                _r_outputs = req_to_outputs.get(rid, [])
+                _num_tokens = count_tokens_from_outputs(_r_outputs)
+                _agg_total_tokens += _num_tokens
+                _agg_total_gen_time_ms += _gen_ms
+                _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))
+
+            # Emit per-request results
+            for rid in batch_request_ids:
+                r_outputs = req_to_outputs.get(rid, [])
+                try:
+                    use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
+                    _metrics = {
+                        "num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
+                        "stage_gen_time_ms": _gen_ms,
+                        "batch_id": int(_batch_seq),
+                        "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
+                        "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
+                        "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
+                    }
+                    if _stats_file:
+                        compute_and_log_stage_request_stats(
+                            _stats_file,
+                            stage_id,
+                            rid,
+                            len(batch_tasks),
+                            r_outputs,
+                            float(_gen_ms),
+                            int(_metrics["rx_transfer_bytes"]),   # type: ignore[index]
+                            float(_metrics["rx_decode_time_ms"]), # type: ignore[index]
+                        )
+                    if use_shm:
+                        out_q.put({
+                            "request_id": rid,
+                            "stage_id": stage_id,
+                            "engine_outputs_shm": payload,
+                            "metrics": _metrics,
+                        })
+                    else:
+                        out_q.put({
+                            "request_id": rid,
+                            "stage_id": stage_id,
+                            "engine_outputs": payload,
+                            "metrics": _metrics,
+                        })
+                except Exception:
+                    out_q.put({
+                        "request_id": rid,
+                        "stage_id": stage_id,
+                        "engine_outputs": r_outputs,
+                        "metrics": {
+                            "num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
+                            "stage_gen_time_ms": _gen_ms,
+                            "rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
+                            "rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
+                            "rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
+                        },
+                    })
+                _logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid)
+        except Exception as e:
+            _logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e)
+            for rid in batch_request_ids:
+                out_q.put({
+                    "request_id": rid,
+                    "stage_id": stage_id,
+                    "error": str(e),
+                })

vllm_omni/entrypoints/stage_utils.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Tuple, Union
+
+import logging
+import json
+import os
+import pickle
+from multiprocessing import shared_memory as _shm
+
+import cloudpickle
+
+logger = logging.getLogger(__name__)
+
+
+def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> None:
+    """Configure per-stage CUDA visibility and current device.
+
+    Behavior
+    - Comma-separated string (e.g. "2,5,7"): set CUDA_VISIBLE_DEVICES exactly
+      to this list; logical index 0 is used as current device.
+    - Integer or digit-string: treat as logical index (0-based) into the current
+      CUDA_VISIBLE_DEVICES mapping; map to the physical device, and then set
+      CUDA_VISIBLE_DEVICES to this single device.
+    - None/"cpu": keep default visibility.
+    - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string.
+    """
+    try:
+        selected_physical: Optional[int] = None
+        logical_idx: Optional[int] = None
+
+        if isinstance(devices, str) and "," in devices:
+            os.environ["CUDA_VISIBLE_DEVICES"] = devices
+            toks = [t.strip() for t in devices.split(",") if t.strip() != ""]
+            if toks:
+                try:
+                    selected_physical = int(toks[0])
+                    logger.debug(
+                        "[Stage-%s] Set CUDA_VISIBLE_DEVICES to %s; logical 0 -> physical %s",
+                        stage_id, devices, selected_physical,
+                    )
+                except Exception as e:
+                    logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
+                    selected_physical = None
+        elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()):
+            logical_idx = max(0, int(devices))
+            vis = os.environ.get("CUDA_VISIBLE_DEVICES")
+            if vis:
+                try:
+                    mapping = [int(x) for x in vis.split(",") if x.strip() != ""]
+                    if 0 <= logical_idx < len(mapping):
+                        selected_physical = mapping[logical_idx]
+                except Exception as e:
+                    logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e)
+                    selected_physical = None
+            if selected_physical is None:
+                selected_physical = int(logical_idx)
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
+            logger.debug(
+                "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES to single device",
+                stage_id, logical_idx + 1, selected_physical,
+            )
+        elif devices in (None, "cpu"):
+            logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
+        else:
+            selected_physical = int(str(devices))
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
+            logger.debug("[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical)
+
+        try:
+            import torch  # noqa: WPS433
+            if torch.cuda.is_available():
+                try:
+                    torch.cuda.set_device(0)
+                except Exception as e:
+                    logger.debug("[Stage-%s] torch.cuda.set_device(0) failed: %s", stage_id, e, exc_info=True)
+                num = torch.cuda.device_count()
+                info = []
+                for i in range(num):
+                    total = torch.cuda.get_device_properties(i).total_memory
+                    free, _ = torch.cuda.mem_get_info(i)
+                    info.append({
+                        "idx": i,
+                        "name": torch.cuda.get_device_name(i),
+                        "total": int(total),
+                        "free": int(free),
+                    })
+                logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info)
+        except Exception as e:
+            logger.debug("[Stage-%s] Failed to query CUDA devices: %s", stage_id, e, exc_info=True)
+    except Exception as e:
+        logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)
+
+
+def serialize_obj(obj: Any) -> bytes:
+    """Serialize a Python object to bytes using cloudpickle."""
+    return cloudpickle.dumps(obj)
+
+
+def shm_write_bytes(payload: bytes) -> Dict[str, Any]:
+    """Write bytes into SharedMemory and return meta dict {name,size}.
+
+    Caller should close the segment; the receiver should unlink.
+    """
+    shm = _shm.SharedMemory(create=True, size=len(payload))
+    mv = memoryview(shm.buf)
+    mv[: len(payload)] = payload
+    del mv
+    meta = {"name": shm.name, "size": len(payload)}
+    try:
+        shm.close()
+    except Exception as e:
+        logger.debug("Ensure parent dir failed for %s: %s", path, e)
+    return meta
+
+
+def shm_read_bytes(meta: Dict[str, Any]) -> bytes:
+    """Read bytes from SharedMemory by meta {name,size} and cleanup."""
+    shm = _shm.SharedMemory(name=meta["name"])  # type: ignore[index]
+    mv = memoryview(shm.buf)
+    data = bytes(mv[: meta["size"]])
+    del mv
+    try:
+        shm.close()
+    except Exception:
+        pass
+    try:
+        shm.unlink()
+    except Exception:
+        pass
+    return data
+
+
+def _ensure_parent_dir(path: str) -> None:
+    """Ensure the parent directory for a file path exists (best-effort)."""
+    try:
+        parent = os.path.dirname(path)
+        if parent:
+            os.makedirs(parent, exist_ok=True)
+    except Exception:
+        pass
+
+
+def append_jsonl(path: str, record: Dict[str, Any]) -> None:
+    """Append a JSON record as one line to a JSONL file (best-effort).
+
+    This is safe to call from multiple processes when each process writes
+    to a distinct file. For concurrent writes to the same file, OS append
+    semantics typically suffice, but no additional locking is provided.
+    """
+    try:
+        _ensure_parent_dir(path)
+        line = json.dumps(record, ensure_ascii=False)
+        fd = os.open(path, os.O_APPEND | os.O_CREAT | os.O_WRONLY, 0o644)
+        with os.fdopen(fd, "a", encoding="utf-8") as f:
+            f.write(line + "\n")
+    except Exception:
+        logger.exception("Failed to append JSONL to %s", path)
+
+
+def maybe_dump_to_shm(obj: Any, threshold: int) -> Tuple[bool, Any]:
+    """Dump object to SHM if serialized size exceeds threshold.
+
+    Returns (True, meta) when dumped; otherwise (False, original_obj).
+    """
+    payload = serialize_obj(obj)
+    if len(payload) > threshold:
+        return True, shm_write_bytes(payload)
+    return False, obj
+
+
+def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) -> Any:
+    """Load object from container that may carry SHM or inline object.
+
+    Deprecated: prefer `maybe_load_from_ipc_with_metrics` to also obtain
+    decode-time and size metrics.
+    """
+    if shm_key in container:
+        return pickle.loads(shm_read_bytes(container[shm_key]))
+    return container[obj_key]
+
+
+def maybe_load_from_ipc_with_metrics(container: Dict[str, Any], obj_key: str, shm_key: str) -> tuple[Any, Dict[str, float]]:
+    """Load object and return (object, metrics) with RX bytes and decode time.
+
+    Metrics keys:
+      - rx_transfer_bytes: int
+      - rx_decode_time_ms: float
+    """
+    import time as _time  # local import to avoid overhead at module import
+    t0 = _time.time()
+    if shm_key in container:
+        meta = container[shm_key]  # type: ignore[index]
+        payload = shm_read_bytes(meta)
+        obj = pickle.loads(payload)
+        try:
+            rx_bytes = int(meta.get("size", len(payload)))  # type: ignore[call-arg]
+        except Exception:
+            rx_bytes = len(payload)
+    else:
+        obj = container[obj_key]
+        try:
+            rx_bytes = len(serialize_obj(obj))
+        except Exception:
+            rx_bytes = 0
+    t1 = _time.time()
+    rx_decode_ms = (t1 - t0) * 1000.0
+    return obj, {
+        "rx_transfer_bytes": int(rx_bytes),
+        "rx_decode_time_ms": float(rx_decode_ms),
+    }
+
+
+def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> Dict[str, Any]:
+    """Return a dict payload for IPC: inline (obj_key) or SHM (shm_key).
+
+    When serialized size exceeds threshold, returns {shm_key: {name,size}};
+    otherwise returns {obj_key: obj}.
+    """
+    payload: Dict[str, Any] = {}
+    use_shm, data = maybe_dump_to_shm(obj, threshold)
+    if use_shm:
+        payload[shm_key] = data
+    else:
+        payload[obj_key] = data
+    return payload
+
+
+# Convert OmegaConf/objects to plain dicts
+def _to_dict(x: Any) -> Dict[str, Any]:
+    try:
+        if isinstance(x, dict):
+            return dict(x)
+        return OmegaConf.to_container(x, resolve=True)  # type: ignore[arg-type]
+    except Exception:
+        try:
+            return dict(x)
+        except Exception:
+            return {}

vllm_omni/entrypoints/utils.py

@@ -1,6 +1,8 @@
 import os
 from pathlib import Path
+
 from omegaconf import OmegaConf
+
 from vllm.transformers_utils.config import get_config
 
 # Get the project root directory (2 levels up from this file)
@@ -22,4 +24,4 @@ def load_stage_configs_from_model(model: str):
 def load_stage_configs_from_yaml(config_path: str):
     """Load stage configs from yaml file."""
     config_data = OmegaConf.load(config_path)
-    return config_data.stage_args
+    return config_data.stage_args

vllm_omni/inputs/data.py

@@ -1,7 +1,9 @@
-from vllm.inputs.data import TokensPrompt
-from typing import Any, NotRequired
+from typing import Any, NotRequired, Optional
+
 import torch
 
+from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt
+
 
 class OmniTokensPrompt(TokensPrompt):
     prompt_embeds: NotRequired[torch.Tensor]
@@ -9,4 +11,45 @@ class OmniTokensPrompt(TokensPrompt):
 
     # New: optional additional information dictionary
     # Values may be torch.Tensor or list
-    additional_information: NotRequired[dict[str, Any]]
+    additional_information: NotRequired[dict[str, Any]]
+
+
+class OmniTokenInputs(TokenInputs):
+    # New: optional prompt embeddings aligned with token ids
+    prompt_embeds: NotRequired[torch.Tensor]
+
+    # New: optional additional information dictionary
+    # Values may be torch.Tensor or list
+    additional_information: NotRequired[dict[str, Any]]
+
+
+class OmniEmbedsPrompt(EmbedsPrompt):
+    # New: optional prompt embeddings aligned with token ids
+    prompt_embeds: NotRequired[torch.Tensor]
+
+    # New: optional additional information dictionary
+    # Values may be torch.Tensor or list
+    additional_information: NotRequired[dict[str, Any]]
+
+
+def token_inputs_omni(
+    prompt_token_ids: list[int],
+    prompt: Optional[str] = None,
+    cache_salt: Optional[str] = None,
+    prompt_embeds: Optional[torch.Tensor] = None,
+    additional_information: Optional[dict[str, Any]] = None,
+) -> OmniTokenInputs:
+    """Construct [`TokenInputs`][vllm.inputs.data.TokenInputs] from optional
+    values."""
+    inputs = OmniTokenInputs(type="token", prompt_token_ids=prompt_token_ids)
+
+    if prompt is not None:
+        inputs["prompt"] = prompt
+    if cache_salt is not None:
+        inputs["cache_salt"] = cache_salt
+    if prompt_embeds is not None:
+        inputs["prompt_embeds"] = prompt_embeds
+    if additional_information is not None:
+        inputs["additional_information"] = additional_information
+
+    return inputs

vllm_omni/inputs/parse.py

@@ -0,0 +1,31 @@
+from vllm.inputs.data import SingletonPrompt
+from vllm.inputs.parse import (
+    ParsedEmbedsPrompt,
+    ParsedSingletonPrompt,
+    ParsedStrPrompt,
+    ParsedTextPrompt,
+    ParsedTokensPrompt,
+)
+
+
+def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPrompt:
+    if isinstance(prompt, str):
+        return ParsedStrPrompt(type="str", content=prompt)
+    elif isinstance(prompt, dict):
+        # Type ignores are because mypy does not correctly infer the TypedDicts
+        # Pyright does succeed.
+        # Priority tokens: When both tokens and embeds exist, keep both and
+        # follow the tokens path
+        if "prompt_token_ids" in prompt:
+            return ParsedTokensPrompt(
+                type="tokens", content=prompt
+            )  # type: ignore[typeddict-item]
+        elif "prompt_embeds" in prompt:
+            return ParsedEmbedsPrompt(
+                type="embeds", content=prompt
+            )  # type: ignore[typeddict-item]
+        elif "prompt" in prompt:
+            return ParsedTextPrompt(type="text", content=prompt)
+    raise TypeError(
+        "inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt"
+    )

vllm_omni/inputs/preprocess.py

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

vllm_omni/model_executor/layers/mrope.py

@@ -6,7 +6,7 @@ import torch
 from transformers import PretrainedConfig
 
 from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding
-from vllm.model_executor.layers.rotary_embedding.common import apply_rotary_emb_dispatch
+
 
 def _apply_rotary_emb(
     x: torch.Tensor,
@@ -36,6 +36,7 @@ def _apply_rotary_emb(
     else:
         return torch.stack((o1, o2), dim=-1).flatten(-2)
 
+
 class MRotaryEmbedding(RotaryEmbedding):
     """Rotary Embedding with Multimodal Sections."""
 
@@ -48,17 +49,34 @@ class MRotaryEmbedding(RotaryEmbedding):
         is_neox_style: bool,
         dtype: torch.dtype,
         mrope_section: Optional[list[int]] = None,
+        mrope_interleaved: bool = False,
+        # YaRN parameters.
+        *,
+        scaling_factor: Optional[float] = None,
+        extrapolation_factor: float = 1,
+        attn_factor: float = 1,
+        beta_fast: int = 32,
+        beta_slow: int = 1,
     ) -> None:
         # In Qwen2.5-VL, the maximum index value is related to the duration of
         # the input video. We enlarge max_position_embeddings to 4 times to get
         # a larger the cos and sin cache.
         self.cache_max_position_num = max_position_embeddings * 4
-        super().__init__(head_size, rotary_dim, self.cache_max_position_num,
-                         base, is_neox_style, dtype)
+        super().__init__(
+            head_size,
+            rotary_dim,
+            self.cache_max_position_num,
+            base,
+            is_neox_style,
+            dtype,
+        )
 
         self.mrope_section = mrope_section
         if self.mrope_section:
-            print("Warning: mrope_section check is disabled in Qwen2.5-Omni, this may cause errors, and should be restored in the future.")
+            print(
+                "Warning: mrope_section check is disabled in Qwen2.5-Omni, "
+                "this may cause errors, and should be restored in the future."
+            )
             # assert sum(self.mrope_section) == rotary_dim // 2
 
     def forward(
@@ -85,35 +103,31 @@ class MRotaryEmbedding(RotaryEmbedding):
         if positions.ndim == 2:
             assert self.mrope_section
 
-            cos = torch.cat([
-                m[i]
-                for i, m in enumerate(cos.split(self.mrope_section, dim=-1))
-            ],
-                            dim=-1)
-            sin = torch.cat([
-                m[i]
-                for i, m in enumerate(sin.split(self.mrope_section, dim=-1))
-            ],
-                            dim=-1)
+            cos = torch.cat(
+                [m[i] for i, m in enumerate(cos.split(self.mrope_section, dim=-1))],
+                dim=-1,
+            )
+            sin = torch.cat(
+                [m[i] for i, m in enumerate(sin.split(self.mrope_section, dim=-1))],
+                dim=-1,
+            )
 
         query_shape = query.shape
         query = query.view(num_tokens, -1, self.head_size)
-        query_rot = query[..., :self.rotary_dim]
-        query_pass = query[..., self.rotary_dim:]
+        query_rot = query[..., : self.rotary_dim]
+        query_pass = query[..., self.rotary_dim :]
         # query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin,
         #                                       self.is_neox_style)
-        query_rot = _apply_rotary_emb(query_rot, cos, sin,
-                                              self.is_neox_style)
+        query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
         query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
 
         key_shape = key.shape
         key = key.view(num_tokens, -1, self.head_size)
-        key_rot = key[..., :self.rotary_dim]
-        key_pass = key[..., self.rotary_dim:]
+        key_rot = key[..., : self.rotary_dim]
+        key_pass = key[..., self.rotary_dim :]
         # key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin,
         #                                     self.is_neox_style)
-        key_rot = _apply_rotary_emb(key_rot, cos, sin,
-                                              self.is_neox_style)
+        key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
         key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
         return query, key
 
@@ -134,21 +148,19 @@ class MRotaryEmbedding(RotaryEmbedding):
 
         image_grid_thw = [] if image_grid_thw is None else image_grid_thw
         video_grid_thw = [] if video_grid_thw is None else video_grid_thw
-        second_per_grid_ts = [] if second_per_grid_ts is None else \
-            second_per_grid_ts
+        second_per_grid_ts = [] if second_per_grid_ts is None else second_per_grid_ts
 
-        llm_positions, mrope_position_delta = \
-            cls.get_input_positions_tensor(
-                input_tokens=input_tokens,
-                hf_config=hf_config,
-                image_grid_thw=image_grid_thw,
-                video_grid_thw=video_grid_thw,
-                second_per_grid_ts=second_per_grid_ts,
-                context_len=context_len,
-                seq_len=seq_len,
-                audio_feature_lengths=audio_feature_lengths,
-                use_audio_in_video=use_audio_in_video,
-            )
+        llm_positions, mrope_position_delta = cls.get_input_positions_tensor(
+            input_tokens=input_tokens,
+            hf_config=hf_config,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            second_per_grid_ts=second_per_grid_ts,
+            context_len=context_len,
+            seq_len=seq_len,
+            audio_feature_lengths=audio_feature_lengths,
+            use_audio_in_video=use_audio_in_video,
+        )
 
         return llm_positions.tolist(), mrope_position_delta
 
@@ -166,6 +178,7 @@ class MRotaryEmbedding(RotaryEmbedding):
         use_audio_in_video: bool = False,
     ) -> tuple[torch.Tensor, int]:
         from vllm.transformers_utils.config import thinker_uses_mrope
+
         if thinker_uses_mrope(hf_config):
             return cls._omni_get_input_positions_tensor(
                 input_tokens=input_tokens,
@@ -237,7 +250,8 @@ class MRotaryEmbedding(RotaryEmbedding):
 
             input_type_group: list[tuple[str, int, int]] = []
             for key, group_iter in itertools.groupby(
-                    enumerate(input_token_type), lambda x: x[1]):
+                enumerate(input_token_type), lambda x: x[1]
+            ):
                 group_list = list(group_iter)
                 start_index = group_list[0][0]
                 end_index = group_list[-1][0] + 1
@@ -246,25 +260,42 @@ class MRotaryEmbedding(RotaryEmbedding):
             video_frame_num = 1
             mm_data_idx = 0
             for modality_type, start_idx, end_idx in input_type_group:
-                st_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                    llm_pos_ids_list) > 0 else 0
+                st_idx = (
+                    llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                )
                 if modality_type == "image":
                     t, h, w = (
                         image_grid_thw[mm_data_idx][0],
                         image_grid_thw[mm_data_idx][1],
                         image_grid_thw[mm_data_idx][2],
                     )
-                    llm_grid_t, llm_grid_h, llm_grid_w = \
-                        t, h // spatial_merge_size, w // spatial_merge_size
+                    llm_grid_t, llm_grid_h, llm_grid_w = (
+                        t,
+                        h // spatial_merge_size,
+                        w // spatial_merge_size,
+                    )
 
-                    t_index = torch.arange(llm_grid_t).view(-1, 1).expand(
-                        -1, llm_grid_h * llm_grid_w).flatten()
-                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
-                        llm_grid_t, -1, llm_grid_w).flatten()
-                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
-                        llm_grid_t, llm_grid_h, -1).flatten()
+                    t_index = (
+                        torch.arange(llm_grid_t)
+                        .view(-1, 1)
+                        .expand(-1, llm_grid_h * llm_grid_w)
+                        .flatten()
+                    )
+                    h_index = (
+                        torch.arange(llm_grid_h)
+                        .view(1, -1, 1)
+                        .expand(llm_grid_t, -1, llm_grid_w)
+                        .flatten()
+                    )
+                    w_index = (
+                        torch.arange(llm_grid_w)
+                        .view(1, 1, -1)
+                        .expand(llm_grid_t, llm_grid_h, -1)
+                        .flatten()
+                    )
                     llm_pos_ids_list.append(
-                        torch.stack([t_index, h_index, w_index]) + st_idx)
+                        torch.stack([t_index, h_index, w_index]) + st_idx
+                    )
                     mm_data_idx += 1
 
                 elif modality_type == "video":
@@ -273,18 +304,34 @@ class MRotaryEmbedding(RotaryEmbedding):
                         image_grid_thw[mm_data_idx][1],
                         image_grid_thw[mm_data_idx][2],
                     )
-                    llm_grid_t, llm_grid_h, llm_grid_w = \
-                        t, h // spatial_merge_size, w // spatial_merge_size
+                    llm_grid_t, llm_grid_h, llm_grid_w = (
+                        t,
+                        h // spatial_merge_size,
+                        w // spatial_merge_size,
+                    )
 
                     for t_idx in range(llm_grid_t):
-                        t_index = torch.tensor(t_idx).view(-1, 1).expand(
-                            -1, llm_grid_h * llm_grid_w).flatten()
-                        h_index = torch.arange(llm_grid_h).view(
-                            1, -1, 1).expand(1, -1, llm_grid_w).flatten()
-                        w_index = torch.arange(llm_grid_w).view(
-                            1, 1, -1).expand(1, llm_grid_h, -1).flatten()
+                        t_index = (
+                            torch.tensor(t_idx)
+                            .view(-1, 1)
+                            .expand(-1, llm_grid_h * llm_grid_w)
+                            .flatten()
+                        )
+                        h_index = (
+                            torch.arange(llm_grid_h)
+                            .view(1, -1, 1)
+                            .expand(1, -1, llm_grid_w)
+                            .flatten()
+                        )
+                        w_index = (
+                            torch.arange(llm_grid_w)
+                            .view(1, 1, -1)
+                            .expand(1, llm_grid_h, -1)
+                            .flatten()
+                        )
                         llm_pos_ids_list.append(
-                            torch.stack([t_index, h_index, w_index]) + st_idx)
+                            torch.stack([t_index, h_index, w_index]) + st_idx
+                        )
 
                     mm_data_idx += 1
                     video_frame_num += 1
@@ -292,19 +339,17 @@ class MRotaryEmbedding(RotaryEmbedding):
                 else:
                     text_len = end_idx - start_idx
                     llm_pos_ids_list.append(
-                        torch.arange(text_len).view(1, -1).expand(3, -1) +
-                        st_idx)
+                        torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
+                    )
                     video_frame_num = 1
 
         else:
             text_len = len(input_tokens)
-            llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1))
+            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1))
 
         llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
         llm_positions = llm_positions[:, context_len:seq_len]
-        mrope_position_delta = (llm_positions.max() + 1 -
-                                len(input_tokens)).item()
+        mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
         return llm_positions, mrope_position_delta
 
     @classmethod
@@ -324,12 +369,12 @@ class MRotaryEmbedding(RotaryEmbedding):
         video_token_id = hf_config.video_token_id
         vision_start_token_id = hf_config.vision_start_token_id
         spatial_merge_size = hf_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(hf_config.vision_config,
-                                    "tokens_per_second", 1.0)
+        tokens_per_second = getattr(hf_config.vision_config, "tokens_per_second", 1.0)
 
         input_tokens_tensor = torch.tensor(input_tokens)
         vision_start_indices = torch.argwhere(
-            input_tokens_tensor == vision_start_token_id).squeeze(1)
+            input_tokens_tensor == vision_start_token_id
+        ).squeeze(1)
         vision_tokens = input_tokens_tensor[vision_start_indices + 1]
         image_nums = (vision_tokens == image_token_id).sum()
         video_nums = (vision_tokens == video_token_id).sum()
@@ -371,37 +416,56 @@ class MRotaryEmbedding(RotaryEmbedding):
                 remain_videos -= 1
                 ed = ed_video
 
-            llm_grid_t, llm_grid_h, llm_grid_w = \
-                t, h // spatial_merge_size, w // spatial_merge_size
+            llm_grid_t, llm_grid_h, llm_grid_w = (
+                t,
+                h // spatial_merge_size,
+                w // spatial_merge_size,
+            )
             text_len = ed - st
 
-            st_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                llm_pos_ids_list) > 0 else 0
+            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
             llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
+            )
 
-            t_index = (torch.arange(llm_grid_t).view(-1, 1).expand(
-                -1, llm_grid_h * llm_grid_w) * video_second_per_grid_t *
-                       tokens_per_second).long().flatten()
+            t_index = (
+                (
+                    torch.arange(llm_grid_t)
+                    .view(-1, 1)
+                    .expand(-1, llm_grid_h * llm_grid_w)
+                    * video_second_per_grid_t
+                    * tokens_per_second
+                )
+                .long()
+                .flatten()
+            )
 
-            h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
-                llm_grid_t, -1, llm_grid_w).flatten()
-            w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
-                llm_grid_t, llm_grid_h, -1).flatten()
+            h_index = (
+                torch.arange(llm_grid_h)
+                .view(1, -1, 1)
+                .expand(llm_grid_t, -1, llm_grid_w)
+                .flatten()
+            )
+            w_index = (
+                torch.arange(llm_grid_w)
+                .view(1, 1, -1)
+                .expand(llm_grid_t, llm_grid_h, -1)
+                .flatten()
+            )
             llm_pos_ids_list.append(
-                torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+                torch.stack([t_index, h_index, w_index]) + text_len + st_idx
+            )
             st = ed + llm_grid_t * llm_grid_h * llm_grid_w
 
         if st < len(input_tokens):
-            st_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                llm_pos_ids_list) > 0 else 0
+            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
             text_len = len(input_tokens) - st
             llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
+            )
 
         llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
-        mrope_position_delta = (llm_positions.max() + 1 -
-                                len(input_tokens)).item()
+        mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
         llm_positions = llm_positions[:, context_len:seq_len]
 
         return llm_positions, mrope_position_delta
@@ -448,8 +512,9 @@ class MRotaryEmbedding(RotaryEmbedding):
         vision_end_token_id = thinker_config.vision_end_token_id
         seconds_per_chunk = thinker_config.seconds_per_chunk
         spatial_merge_size = thinker_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(thinker_config.vision_config,
-                                    "tokens_per_second", 25)
+        tokens_per_second = getattr(
+            thinker_config.vision_config, "tokens_per_second", 25
+        )
 
         if isinstance(image_grid_thw, list):
             image_grid_thw = torch.tensor(image_grid_thw)
@@ -469,28 +534,30 @@ class MRotaryEmbedding(RotaryEmbedding):
         idx = 0
         while idx < len(src_item):
             new_src_item_len = len(new_src_item)
-            start_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                llm_pos_ids_list) > 0 else 0
-            if src_item[idx] not in [
-                    audio_token_id, video_token_id, image_token_id
-            ]:
+            start_idx = (
+                llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+            )
+            if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]:
                 if use_audio_in_video and idx > 0:
-                    if src_item[idx] == vision_end_token_id and \
-                        src_item[idx - 1] == audio_end_token_id:
+                    if (
+                        src_item[idx] == vision_end_token_id
+                        and src_item[idx - 1] == audio_end_token_id
+                    ):
                         # processing the <|audio_eos|> before <|vision_eos|>
                         start_idx -= 1
-                    elif src_item[idx] == audio_start_token_id and \
-                        src_item[idx - 1] == vision_start_token_id:
+                    elif (
+                        src_item[idx] == audio_start_token_id
+                        and src_item[idx - 1] == vision_start_token_id
+                    ):
                         # processing the <|audio_bos|> after <|vision_eos|>
                         start_idx -= 1
                 new_src_item.append(src_item[idx])
-                llm_pos_ids = torch.tensor([start_idx],
-                                           dtype=torch.long).expand(3, -1)
+                llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1)
                 llm_pos_ids_list.append(llm_pos_ids)
             elif src_item[idx] == audio_token_id:
                 assert audio_seqlens is not None
                 audio_seqlen = audio_seqlens[audio_idx]
-                place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1)
+                place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1
                 new_src_item.extend([audio_token_id] * place_num)
                 llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
                 llm_pos_ids_list.append(llm_pos_ids)
@@ -501,26 +568,30 @@ class MRotaryEmbedding(RotaryEmbedding):
                 grid_ws = image_grid_thw[:, 2]
                 t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
                 llm_pos_ids = cls._get_llm_pos_ids_for_vision(
-                    start_idx, image_idx, spatial_merge_size, t_index, grid_hs,
-                    grid_ws)
+                    start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+                )
                 llm_pos_ids_list.append(llm_pos_ids)
                 vision_seqlen = image_grid_thw[image_idx].prod() // (
-                    spatial_merge_size**2)
+                    spatial_merge_size**2
+                )
                 new_src_item.extend([image_token_id] * vision_seqlen)
                 image_idx += 1
             elif src_item[idx] == video_token_id and not use_audio_in_video:
                 grid_t = video_grid_thw[video_idx][0]
                 grid_hs = video_grid_thw[:, 1]
                 grid_ws = video_grid_thw[:, 2]
-                t_index = (torch.arange(grid_t) *
-                           second_per_grid_ts[video_idx] *
-                           tokens_per_second).long()
+                t_index = (
+                    torch.arange(grid_t)
+                    * second_per_grid_ts[video_idx]
+                    * tokens_per_second
+                ).long()
                 llm_pos_ids = cls._get_llm_pos_ids_for_vision(
-                    start_idx, video_idx, spatial_merge_size, t_index, grid_hs,
-                    grid_ws)
+                    start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+                )
                 llm_pos_ids_list.append(llm_pos_ids)
                 vision_seqlen = video_grid_thw[video_idx].prod() // (
-                    spatial_merge_size**2)
+                    spatial_merge_size**2
+                )
                 new_src_item.extend([video_token_id] * vision_seqlen)
                 video_idx += 1
             else:
@@ -528,56 +599,73 @@ class MRotaryEmbedding(RotaryEmbedding):
                 assert audio_seqlens is not None
                 audio_seqlen = audio_seqlens[audio_idx]
                 vision_seqlen = video_grid_thw[video_idx].prod() // (
-                    spatial_merge_size**2)
+                    spatial_merge_size**2
+                )
                 grid_t = video_grid_thw[video_idx][0]
                 grid_h = video_grid_thw[video_idx][1]
                 grid_w = video_grid_thw[video_idx][2]
                 grid_hs = video_grid_thw[:, 1]
                 grid_ws = video_grid_thw[:, 2]
                 t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
-                t_index = (torch.arange(grid_t) *
-                           second_per_grid_ts[video_idx] *
-                           tokens_per_second).long()
+                t_index = (
+                    torch.arange(grid_t)
+                    * second_per_grid_ts[video_idx]
+                    * tokens_per_second
+                ).long()
                 t_index_split_chunk = cls._split_list_into_ranges(
-                    t_index, t_ntoken_per_chunk)
+                    t_index, t_ntoken_per_chunk
+                )
                 place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
                 pure_audio_len = place_num - 2
                 added_audio_len = 0
                 audio_llm_pos_ids_list: list[torch.Tensor] = []
                 for t_chunk in t_index_split_chunk:
-                    vision_ntoken_per_chunk = len(
-                        t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
-                    new_src_item.extend([video_token_id] *
-                                        vision_ntoken_per_chunk)
+                    vision_ntoken_per_chunk = (
+                        len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
+                    )
+                    new_src_item.extend([video_token_id] * vision_ntoken_per_chunk)
                     vision_llm_pos_ids_list = cls._get_llm_pos_ids_for_vision(
-                        start_idx, video_idx, spatial_merge_size, t_chunk,
-                        grid_hs, grid_ws).split(1, dim=1)
+                        start_idx,
+                        video_idx,
+                        spatial_merge_size,
+                        t_chunk,
+                        grid_hs,
+                        grid_ws,
+                    ).split(1, dim=1)
                     llm_pos_ids_list.extend(vision_llm_pos_ids_list)
                     new_src_item.extend(
-                        min(t_ntoken_per_chunk, pure_audio_len -
-                            added_audio_len) * [audio_token_id])
-                    audio_start_idx = start_idx if len(
-                        audio_llm_pos_ids_list
-                    ) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
-                    if min(t_ntoken_per_chunk,
-                           pure_audio_len - added_audio_len) > 0:
-                        audio_llm_pos_ids_list = (torch.arange(
-                            min(t_ntoken_per_chunk, pure_audio_len -
-                                added_audio_len)).expand(3, -1) +
-                                                  audio_start_idx).split(1,
-                                                                         dim=1)
+                        min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
+                        * [audio_token_id]
+                    )
+                    audio_start_idx = (
+                        start_idx
+                        if len(audio_llm_pos_ids_list) == 0
+                        else audio_llm_pos_ids_list[-1][0].item() + 1
+                    )
+                    if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0:
+                        audio_llm_pos_ids_list = (
+                            torch.arange(
+                                min(
+                                    t_ntoken_per_chunk, pure_audio_len - added_audio_len
+                                )
+                            ).expand(3, -1)
+                            + audio_start_idx
+                        ).split(1, dim=1)
                     else:
                         audio_llm_pos_ids_list = []
-                    added_audio_len += min(t_ntoken_per_chunk,
-                                           pure_audio_len - added_audio_len)
+                    added_audio_len += min(
+                        t_ntoken_per_chunk, pure_audio_len - added_audio_len
+                    )
                     llm_pos_ids_list.extend(audio_llm_pos_ids_list)
                 if added_audio_len < pure_audio_len:
                     new_src_item.extend(
-                        (pure_audio_len - added_audio_len) * [audio_token_id])
+                        (pure_audio_len - added_audio_len) * [audio_token_id]
+                    )
                     audio_llm_pos_ids_list = (
-                        torch.arange(pure_audio_len - added_audio_len).expand(
-                            3, -1) + llm_pos_ids_list[-1].max() + 1).split(
-                                1, dim=1)
+                        torch.arange(pure_audio_len - added_audio_len).expand(3, -1)
+                        + llm_pos_ids_list[-1].max()
+                        + 1
+                    ).split(1, dim=1)
                     llm_pos_ids_list.extend(audio_llm_pos_ids_list)
                 audio_idx += 1
                 video_idx += 1
@@ -585,8 +673,9 @@ class MRotaryEmbedding(RotaryEmbedding):
             idx += len(new_src_item) - new_src_item_len
 
         llm_positions = torch.cat(llm_pos_ids_list, dim=1)
-        mrope_position_delta = torch.cat(llm_pos_ids_list,
-                                         dim=1).max() + 1 - len(src_item)
+        mrope_position_delta = (
+            torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item)
+        )
         llm_positions = llm_positions[:, context_len:seq_len]
 
         return llm_positions, mrope_position_delta
@@ -603,22 +692,34 @@ class MRotaryEmbedding(RotaryEmbedding):
         llm_pos_ids_list = []
         llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
         llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
-        h_index = (torch.arange(llm_grid_h).view(1, -1, 1).expand(
-            len(t_index), -1, llm_grid_w).flatten())
-        w_index = (torch.arange(llm_grid_w).view(1, 1, -1).expand(
-            len(t_index), llm_grid_h, -1).flatten())
-        t_index_tensor = torch.Tensor(t_index).to(llm_grid_h.device).view(
-            -1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten()
+        h_index = (
+            torch.arange(llm_grid_h)
+            .view(1, -1, 1)
+            .expand(len(t_index), -1, llm_grid_w)
+            .flatten()
+        )
+        w_index = (
+            torch.arange(llm_grid_w)
+            .view(1, 1, -1)
+            .expand(len(t_index), llm_grid_h, -1)
+            .flatten()
+        )
+        t_index_tensor = (
+            torch.Tensor(t_index)
+            .to(llm_grid_h.device)
+            .view(-1, 1)
+            .expand(-1, llm_grid_h * llm_grid_w)
+            .long()
+            .flatten()
+        )
         _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
         llm_pos_ids_list.append(_llm_pos_ids + start_idx)
         llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
         return llm_pos_ids
 
     @staticmethod
-    def _split_list_into_ranges(lst: torch.Tensor,
-                                interval: int) -> list[list[int]]:
-        ranges: list[list[int]] = [[]
-                                   for _ in range((max(lst) // interval) + 1)]
+    def _split_list_into_ranges(lst: torch.Tensor, interval: int) -> list[list[int]]:
+        ranges: list[list[int]] = [[] for _ in range((max(lst) // interval) + 1)]
         for num in lst:
             index = num // interval
             ranges[index].append(num)
@@ -632,19 +733,28 @@ class MRotaryEmbedding(RotaryEmbedding):
     ) -> list[list[int]]:
         return [
             list(
-                range(context_len + mrope_position_delta,
-                      seq_len + mrope_position_delta)) for _ in range(3)
+                range(
+                    context_len + mrope_position_delta, seq_len + mrope_position_delta
+                )
+            )
+            for _ in range(3)
         ]
 
     @staticmethod
-    def get_next_input_positions_tensor(out: np.ndarray, out_offset: int,
-                                        mrope_position_delta: int,
-                                        context_len: int, num_new_tokens: int):
+    def get_next_input_positions_tensor(
+        out: np.ndarray,
+        out_offset: int,
+        mrope_position_delta: int,
+        context_len: int,
+        num_new_tokens: int,
+    ):
 
-        values = np.arange(mrope_position_delta + context_len,
-                           mrope_position_delta + context_len + num_new_tokens,
-                           dtype=out.dtype)
-        out[:, out_offset:out_offset + num_new_tokens] = values
+        values = np.arange(
+            mrope_position_delta + context_len,
+            mrope_position_delta + context_len + num_new_tokens,
+            dtype=out.dtype,
+        )
+        out[:, out_offset : out_offset + num_new_tokens] = values
 
     @classmethod
     def omni_get_updates_use_audio_in_video(
@@ -669,27 +779,28 @@ class MRotaryEmbedding(RotaryEmbedding):
         audio_end_token_id = thinker_config.audio_end_token_id
         seconds_per_chunk = thinker_config.seconds_per_chunk
         spatial_merge_size = thinker_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(thinker_config.vision_config,
-                                    "tokens_per_second", 25)
+        tokens_per_second = getattr(
+            thinker_config.vision_config, "tokens_per_second", 25
+        )
 
         grid_t = video_grid_thw[0]
         grid_h = video_grid_thw[1]
         grid_w = video_grid_thw[2]
         t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
-        t_index = (torch.arange(grid_t) * video_second_per_grid_t *
-                   tokens_per_second).long()
-        t_index_split_chunk = cls._split_list_into_ranges(
-            t_index, t_ntoken_per_chunk)
+        t_index = (
+            torch.arange(grid_t) * video_second_per_grid_t * tokens_per_second
+        ).long()
+        t_index_split_chunk = cls._split_list_into_ranges(t_index, t_ntoken_per_chunk)
 
         updates = [audio_start_token_id]
         added_audio_len = 0
         for t_chunk in t_index_split_chunk:
-            vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (
-                spatial_merge_size**2)
+            vision_ntoken_per_chunk = (
+                len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
+            )
             updates.extend([video_token_id] * vision_ntoken_per_chunk)
 
-            audio_chunk_size = min(t_ntoken_per_chunk,
-                                   audio_len - added_audio_len)
+            audio_chunk_size = min(t_ntoken_per_chunk, audio_len - added_audio_len)
             updates.extend(audio_chunk_size * [audio_token_id])
             added_audio_len += audio_chunk_size
         if added_audio_len < audio_len:

vllm_omni/model_executor/model_loader/weight_utils.py

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

vllm_omni/model_executor/models/__init__.py

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

vllm_omni/model_executor/models/qwen2_5_omni.py

@@ -1,77 +1,82 @@
-from functools import cached_property
-from typing import Iterable, List, Optional, Set, Tuple, Union, NamedTuple, Dict
-
-import os
 import glob
+import os
+from functools import cached_property
+from typing import Dict, Iterable, NamedTuple, Optional, Set, Tuple, Union
+
 import numpy as np
 import torch
 import torch.nn as nn
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
-    Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig, Qwen2_5OmniTalkerConfig)
+    Qwen2_5OmniConfig,
+    Qwen2_5OmniTalkerConfig,
+    Qwen2_5OmniThinkerConfig,
+)
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import ColumnParallelLinear
-from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
-from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
     Qwen2_5OmniThinkerProcessingInfo,
-    Qwen2_5OmniThinkerDummyInputsBuilder)
+)
+from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
+
 # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.sequence import IntermediateTensors
-
-from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
-from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
-                    init_vllm_registered_model,
-                    maybe_prefix)
-from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
 
+
 class OmniOutput(NamedTuple):
     """Output from the merged Omni model containing both text and audio."""
+
     text_hidden_states: torch.Tensor
     multimodal_outputs: dict = {}
     intermediate_tensors: Optional[IntermediateTensors] = None
 
+
 logger = init_logger(__name__)
 
+
 @MULTIMODAL_REGISTRY.register_processor(
     Qwen2_5OmniThinkerMultiModalProcessor,
     info=Qwen2_5OmniThinkerProcessingInfo,
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
-class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
-                                         SupportsPP,
-                                         Qwen2_5OmniConditionalGenerationMixin):
+class Qwen2_5OmniForConditionalGeneration(
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+):
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        self.have_multimodal_outputs = True 
+        self.have_multimodal_outputs = True
         config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
         multimodal_config = vllm_config.model_config.multimodal_config
         # keep vllm_config for later submodule init
         self.vllm_config = vllm_config
-        
+
         # Initialize thinker components
         thinker_config: Qwen2_5OmniThinkerConfig = config.thinker_config
         self.thinker_config = thinker_config
         self.multimodal_config = multimodal_config
-        
+
         # Initialize talker components
         talker_config: Qwen2_5OmniTalkerConfig = config.talker_config
         self.talker_config = talker_config
-        
-
 
         self.model_stage = vllm_config.model_config.model_stage
-        if self.model_stage=="thinker":
+        if self.model_stage == "thinker":
             # Initialize thinker model (multimodal processing)
             self.thinker = init_vllm_registered_model(
                 vllm_config=vllm_config,
@@ -79,12 +84,12 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 hf_config=thinker_config,
                 # Use registry architecture key
                 architectures=["Qwen2_5OmniThinkerModel"],
-                )
+            )
             self.model = self.thinker
             self.talker = None
             self.token2wav = None
-            
-        elif self.model_stage=="talker":
+
+        elif self.model_stage == "talker":
             self.thinker = None
             # Initialize talker model wrapper (handles projection + LM)
             self.talker = init_vllm_registered_model(
@@ -95,14 +100,14 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 architectures=["Qwen2_5OmniTalkerModel"],
             )
             self.talker.init_multi_modal(thinker_config)
-            self.model=self.talker
+            self.model = self.talker
             self.token2wav = None
-            
-        elif self.model_stage=="code2wav":
+
+        elif self.model_stage == "code2wav":
             self.thinker = None
             self.talker = None
             # Initialize token2wav (code->mel->wav) like thinker/talker
-            self.token2wav_config = getattr(config, 'token2wav_config', None)
+            self.token2wav_config = getattr(config, "token2wav_config", None)
             self.token2wav = None
             if self.token2wav_config is not None:
                 self.token2wav = init_vllm_registered_model(
@@ -117,10 +122,13 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
             self.model = self.token2wav
         else:
             raise ValueError("Invalid model stage")
-        
+
         # Set up intermediate tensors
         self.make_empty_intermediate_tensors = (
-            self.thinker.make_empty_intermediate_tensors) if self.model_stage=="thinker" else lambda: None
+            (self.thinker.make_empty_intermediate_tensors)
+            if self.model_stage == "thinker"
+            else lambda: None
+        )
 
     # -------------------- Device utilities --------------------
     @staticmethod
@@ -160,7 +168,7 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
     def sampler(self):
         if hasattr(self.model, "sampler"):
             return self.model.sampler
-        return get_sampler()
+        return Sampler()
 
     def get_input_embeddings(
         self,
@@ -168,15 +176,17 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         multimodal_embeddings=None,
     ) -> torch.Tensor:
         if self.model_stage == "code2wav":
-            return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size())
-        return self.model.get_input_embeddings(
-            input_ids, multimodal_embeddings)
-       
+            return (
+                torch.zeros_like(input_ids)
+                .reshape(-1, 1)
+                .repeat(1, self.vllm_config.model_config.get_hidden_size())
+            )
+        return self.model.get_input_embeddings(input_ids, multimodal_embeddings)
 
     def get_multimodal_embeddings(self, **kwargs):
         # Delegate to thinker model for multimodal processing
         return self.model.get_multimodal_embeddings(**kwargs)
-    
+
     def last_index_of(self, list, value):
         return len(list) - 1 - list[::-1].index(value)
 
@@ -191,7 +201,7 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         codec: Optional[torch.Tensor] = None,
         sampling_metadata: Optional[SamplingMetadata] = None,
         logits_index: Optional[int] = None,
-        sampler = None,
+        sampler=None,
         additional_information: Optional[dict[str, object]] = None,
         **kwargs: object,
     ) -> Union[torch.Tensor, IntermediateTensors, OmniOutput]:
@@ -201,8 +211,8 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         2) If audio requested and codec not provided, use talker to derive codec.
         3) If audio requested (or codec provided), use token2wav to synthesize waveform.
         4) Return text hidden states (and audio when applicable).
-        """        
-        if self.model_stage=="thinker":
+        """
+        if self.model_stage == "thinker":
             # Normalize to batched inputs if caller provides 1D/2D unbatched tensors
             added_batch_dim = False
             if input_ids is not None and input_ids.ndim == 1:
@@ -215,10 +225,15 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 inputs_embeds = inputs_embeds.unsqueeze(0)
                 added_batch_dim = True
             thinker_dev = self._module_device(self.thinker)
-            
-            #if input_ids is None, set it to a zero tensor, in the length of the same as the embedding seq length
+
+            # if input_ids is None, set it to a zero tensor, in the length of the
+            # same as the embedding seq length
             if input_ids is None:
-                input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev).unsqueeze(0) #(1, 0)
+                input_ids = torch.zeros(
+                    inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev
+                ).unsqueeze(
+                    0
+                )  # (1, 0)
                 added_batch_dim = True
 
             # 1) Thinker (ensure inputs on thinker's device)
@@ -244,91 +259,201 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
 
             # Text-only path
             return OmniOutput(
-                text_hidden_states=text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states,
-                multimodal_outputs=None
+                text_hidden_states=(
+                    text_hidden_states.squeeze(0)
+                    if added_batch_dim
+                    else text_hidden_states
+                ),
+                multimodal_outputs=None,
             )
 
         # 2) Talker (if codec not provided)
-        if self.model_stage=="talker":
-            
+        if self.model_stage == "talker":
+            # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
+            # Rules:
+            # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
+            # - Decode segments consist of a single non-special token.
+            # - If additional_information is provided (can be a list split by request or a concatenated tensor plus a list of shapes),
+            #   then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments;
+            # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them; otherwise, keep the original embedding.
+
             if input_ids is None and additional_information is None:
-                input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
+                input_ids = torch.zeros(
+                    inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device
+                )
                 additional_information = {}
                 self.thinker_reply_part = torch.zeros_like(inputs_embeds)
                 is_profile = True
             else:
                 is_profile = False
 
-            if input_ids is not None and additional_information is not None and not is_profile:
-                # read from additional_information dict
-                thinker_result = None
-                if additional_information is not None and isinstance(additional_information, dict):
-                    thinker_result = additional_information.get("thinker_result")
-                    prompt_embeds = additional_information.get("prompt_embeds")
-                    prompt_token_ids = additional_information.get("prompt_token_ids")
-                    thinker_output_token_ids = additional_information.get("thinker_output_token_ids")
-                else:
-                    thinker_result = torch.zeros_like(inputs_embeds)
-                    prompt_embeds = torch.zeros_like(inputs_embeds)
-                    prompt_token_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.int64, device=inputs_embeds.device)
-                    thinker_output_token_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.int64, device=inputs_embeds.device)
-                    
-                if thinker_result is None:
-                    thinker_result = torch.zeros_like(inputs_embeds)
-                self.thinker_reply_part = thinker_result.squeeze(0)
-                if self.thinker_reply_part.shape[1] > 1:
-                    self.thinker_reply_part = self.thinker_reply_part[1:, :]
-                input_ids, inputs_embeds = self._thinker_to_talker_prefill(
-                    voice_type=voice_type,
-                    output_prompt_embeds=thinker_result,
-                    output_token_ids = thinker_output_token_ids,
-                    thinker_prompt_embeds=prompt_embeds, 
-                    prompt_token_ids = prompt_token_ids,
-                )
-            elif not is_profile:
-                input_ids, inputs_embeds = \
-                self._thinker_to_talker_decode_one_step(
-                    output_prompt_embeds=self.thinker_reply_part[:1] if self.thinker_reply_part.shape[0]>=1 \
-                        else torch.zeros(1, self.thinker_reply_part.shape[1]).to(self._module_device(self.model)).to(torch.bfloat16)+(-1.25*2**(-123)),
-                    output_token_ids=input_ids,
-                )
+            # Ensure we have base embeddings when only ids are provided
+            if inputs_embeds is None and input_ids is not None:
+                inputs_embeds = self.talker.get_input_embeddings(input_ids)
 
-                if self.thinker_reply_part.shape[0] >=1:
-                    self.thinker_reply_part = self.thinker_reply_part[1:, :]
+            # ------- Request-scoped additional information (no cross-request concat) -------
+            request_ids: Optional[list[str]] = kwargs.get("request_ids")  # ordered
+            request_token_spans: Optional[list[tuple[int,int]]] = kwargs.get("request_token_spans")
+            addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id")
+            runtime_addi = kwargs.get("runtime_additional_information")
+
+            # Normalize runtime_addi into a mapping by request_id for convenience
+            runtime_addi_by_req: dict[str, dict] = {}
+            if isinstance(request_ids, list) and isinstance(runtime_addi, list) and len(runtime_addi) == len(request_ids):
+                for i, rid in enumerate(request_ids):
+                    if isinstance(rid, str) and isinstance(runtime_addi[i], dict):
+                        runtime_addi_by_req[rid] = runtime_addi[i]
+            elif isinstance(request_ids, list) and isinstance(runtime_addi, dict):
+                for rid in request_ids:
+                    if isinstance(rid, str) and isinstance(runtime_addi.get(rid), dict):
+                        runtime_addi_by_req[rid] = runtime_addi[rid]
+
+            # Containers to return per-request updates (e.g., thinker_reply_part_per_request)
+            update_by_req_id: dict[str, dict] = {}
+
+            # ------- Prefill: span_len > 1 -------
+            if (
+                not is_profile
+                and isinstance(request_ids, list)
+                and isinstance(request_token_spans, list)
+                and isinstance(addi_by_req, dict)
+            ):
+                for idx_req, rid in enumerate(request_ids):
+                    s, e = request_token_spans[idx_req]
+                    span_len = int(e) - int(s)
+                    if span_len <= 1:
+                        continue
+                    info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {}
+                    if not isinstance(info, dict):
+                        info = {}
+                    pe = info.get("prompt_embeds")  # Tensor [P,H]
+                    tr = info.get("thinker_result")  # Tensor [K,H]
+                    ptoks = info.get("prompt_token_ids")  # list[int]
+                    otoks = info.get("thinker_output_token_ids")  # list[int]
+
+                    if not isinstance(pe, torch.Tensor):
+                        pe = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
+                    if not isinstance(tr, torch.Tensor):
+                        tr = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
+                    if not isinstance(ptoks, (list, torch.Tensor)):
+                        ptoks = []
+                    if not isinstance(otoks, (list, torch.Tensor)):
+                        otoks = []
+
+                    req_input_ids, req_embeds = self._thinker_to_talker_prefill(
+                        voice_type=voice_type,
+                        output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                        output_token_ids=otoks,
+                        thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                        prompt_token_ids=ptoks,
+                    )
+                    seg_len = min(span_len, req_embeds.shape[0])
+                    inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
+                    if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
+                        input_ids[s : s + seg_len] = req_input_ids
+
+                    # Prepare per-request reply queue for subsequent decode: drop first row
+                    if tr.ndim == 2 and tr.shape[0] > 0:
+                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = tr[1:].detach().to("cpu").contiguous()
+
+            # ------- Decode: span_len == 1 -------
+            if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list):
+                for idx_req, rid in enumerate(request_ids):
+                    s, e = request_token_spans[idx_req]
+                    if (int(e) - int(s)) != 1:
+                        continue
+                    # choose step vector in priority order
+                    step_vec = None
+                    # A) runtime queue
+                    q = None
+                    if isinstance(rid, str):
+                        q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request")
+                    if isinstance(q, torch.Tensor) and q.numel() > 0:
+                        step_vec = q[0:1]
+                        new_q = q[1:].detach().to("cpu").contiguous()
+                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q
+                    else:
+                        # B) per-request provided decode vector (optional)
+                        info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {}
+                        dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None
+                        if isinstance(dv, torch.Tensor) and dv.numel() > 0:
+                            step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
+                        elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0:
+                            # C) fallback shared pool
+                            step_vec = self.thinker_reply_part[0:1]
+                            self.thinker_reply_part = self.thinker_reply_part[1:]
+
+                    if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
+                        one_id = input_ids[s : s + 1]
+                        _, one_embed = self._thinker_to_talker_decode_one_step(
+                            output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                            output_token_ids=one_id,
+                        )
+                        inputs_embeds[s] = one_embed[0]
 
             with torch.inference_mode():
                 talker_hidden = self.talker(
                     input_ids=input_ids,
                     positions=positions[0],
                     inputs_embeds=inputs_embeds,
-
                 )
-            
+            multimodal_outputs: dict = None
+            # Return updates if any
+            if update_by_req_id:
+                multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id}
+
+            if sampling_metadata is not None:
+                # the padding token id is set to text model's pad token id, which do not match with the talker model's word embedding size
+                sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448
+
             return OmniOutput(
                 text_hidden_states=talker_hidden,
-                multimodal_outputs=None
+                multimodal_outputs=multimodal_outputs,
             )
 
-        if self.model_stage=="code2wav":
-            code = input_ids if input_ids is not None else torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
-            audio_tensor = self.generate_audio(code[:-1] if code[-1]==TALKER_CODEC_EOS_TOKEN_ID else code, voice_type)
+        if self.model_stage == "code2wav":
+            code = (
+                input_ids
+                if input_ids is not None
+                else torch.zeros(
+                    inputs_embeds.shape[0],
+                    dtype=torch.long,
+                    device=inputs_embeds.device,
+                )
+            )
+            
+            code = code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code
+            code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code
+            
+            audio_tensor = self.generate_audio(
+                code, voice_type
+            )
             return OmniOutput(
-                text_hidden_states = None,
-                multimodal_outputs = {
-                    "audio": audio_tensor
-                }
+                text_hidden_states=None, multimodal_outputs={"audio": audio_tensor}
             )
 
         return OmniOutput(
             text_hidden_states=torch.cat(
                 [
-                    torch.zeros([inputs_embeds.shape[0],self.talker.config.hidden_size], dtype=torch.bfloat16).to(self._module_device(self.model)),
-                    self.talker.thinker_to_talker_proj(self.talker.get_input_embeddings(torch.tensor([TALKER_CODEC_BOS_TOKEN_ID,TALKER_CODEC_EOS_TOKEN_ID]).to(torch.bfloat16).to(self._module_device(self.model))))[0]
+                    torch.zeros(
+                        [inputs_embeds.shape[0], self.talker.config.hidden_size],
+                        dtype=torch.bfloat16,
+                    ).to(self._module_device(self.model)),
+                    self.talker.thinker_to_talker_proj(
+                        self.talker.get_input_embeddings(
+                            torch.tensor(
+                                [TALKER_CODEC_BOS_TOKEN_ID, TALKER_CODEC_EOS_TOKEN_ID]
+                            )
+                            .to(torch.bfloat16)
+                            .to(self._module_device(self.model))
+                        )
+                    )[0],
                 ],
-                dim=0),
-            multimodal_outputs=None
+                dim=0,
+            ),
+            multimodal_outputs=None,
         )
-    
+
     def generate_audio(self, code, voice_type):
         token2wav_dev = self._module_device(self.token2wav)
         if isinstance(code, torch.Tensor):
@@ -343,7 +468,7 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         return audio_tensor
 
     def _load_talker_embedding(
-            self,
+        self,
     ) -> torch.nn.Embedding:
         return self.talker.language_model.model.embed_tokens
 
@@ -356,19 +481,17 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         # embed_text_bos_token
         self.tts_text_spk_token_ids = {
             # M02: Male voice with standard Mandarin and a slight northern accent
-            'm02': 151870,
-            'Ethan': 151870,
-
+            "m02": 151870,
+            "Ethan": 151870,
             # F030: Your anime-styled virtual girlfriend
-            'f030': 151872,
-            'Chelsie': 151872,
+            "f030": 151872,
+            "Chelsie": 151872,
         }
-        self.default_tts_text_spk_type = list(
-            self.tts_text_spk_token_ids.keys())[0]
-        self.tts_text_spk_token_ids['prefix_caching'] = 151870
+        self.default_tts_text_spk_type = list(self.tts_text_spk_token_ids.keys())[0]
+        self.tts_text_spk_token_ids["prefix_caching"] = 151870
 
         talker_hf_config = self.talker_config
-        if hasattr(talker_hf_config, 'talker_config'):
+        if hasattr(talker_hf_config, "talker_config"):
             talker_hf_config = talker_hf_config.talker_config
 
         self.embed_text_bos_token = self.thinker_embedding(
@@ -376,15 +499,16 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 [talker_hf_config.tts_text_start_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         self.embed_text_spk_tokens = {
-            key:
-            self.thinker_embedding(
+            key: self.thinker_embedding(
                 torch.tensor(
                     [value],
                     dtype=torch.long,
                     device=self._module_device(self.talker),
-                ))
+                )
+            )
             for key, value in self.tts_text_spk_token_ids.items()
         }
         self.embed_text_eos_token = self.thinker_embedding(
@@ -392,31 +516,36 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 [talker_hf_config.tts_text_end_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         self.embed_text_pad_token = self.thinker_embedding(
             torch.tensor(
                 [talker_hf_config.tts_text_pad_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         self.embed_codec_bos_token = self.talker_embedding(
             torch.tensor(
                 [talker_hf_config.tts_codec_start_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         self.embed_codec_eos_token = self.talker_embedding(
             torch.tensor(
                 [talker_hf_config.tts_codec_end_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         self.embed_codec_pad_token = self.talker_embedding(
             torch.tensor(
                 [talker_hf_config.tts_codec_pad_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         return set(["thinker_embedding.weight", "talker_embedding.weight"])
 
     def _get_embed_text_spk_token(self, voice_type: str):
@@ -426,7 +555,7 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
 
     def _get_text_spk_token_id(self, voice_type: str):
         talker_hf_config = self.talker_config
-        if hasattr(talker_hf_config, 'talker_config'):
+        if hasattr(talker_hf_config, "talker_config"):
             talker_hf_config = talker_hf_config.talker_config
 
         if voice_type not in self.tts_text_spk_token_ids:
@@ -439,65 +568,69 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         output_prompt_embeds,
         output_token_ids,
         thinker_prompt_embeds,
-        prompt_token_ids
+        prompt_token_ids,
     ):
 
         talker_hf_config = self.talker_config
-        if hasattr(talker_hf_config, 'talker_config'):
+        if hasattr(talker_hf_config, "talker_config"):
             talker_hf_config = talker_hf_config.talker_config
 
         # if len(output.outputs[0].token_ids) == 2:
-            # issue request
-        prompt_embeds = torch.cat([
-            thinker_prompt_embeds,
-            self._get_embed_text_spk_token(voice_type) +
-            self.embed_codec_pad_token,
-            output_prompt_embeds[:1] + self.embed_codec_bos_token,
-        ],
-                                    dim=0)
+        # issue request
+        prompt_embeds = torch.cat(
+            [
+                thinker_prompt_embeds,
+                self._get_embed_text_spk_token(voice_type) + self.embed_codec_pad_token,
+                output_prompt_embeds[:1] + self.embed_codec_bos_token,
+            ],
+            dim=0,
+        )
 
         prompt_token_ids_processed = prompt_token_ids + [
             talker_hf_config.tts_codec_pad_token_id,
             output_token_ids[0],
-
         ]
         input_tokens_len = len(prompt_token_ids_processed)
         # the code below is from model runner in Qwen, may need to further discuss later
         if input_tokens_len > 2:
-            prompt_token_ids_processed = (
-                [self.talker_config.tts_codec_mask_token_id] *
-                (input_tokens_len - 2) + [
-                    self.talker_config.tts_codec_pad_token_id,
-                    self.talker_config.tts_codec_start_token_id
-                ])
+            prompt_token_ids_processed = [
+                self.talker_config.tts_codec_mask_token_id
+            ] * (input_tokens_len - 2) + [
+                self.talker_config.tts_codec_pad_token_id,
+                self.talker_config.tts_codec_start_token_id,
+            ]
         else:
             prompt_token_ids_processed = [
                 self.talker_config.tts_codec_pad_token_id,
                 self.talker_config.tts_codec_start_token_id,
             ][-input_tokens_len:]
-        if isinstance(prompt_token_ids_processed,list):
-            prompt_token_ids_processed = torch.Tensor(prompt_token_ids_processed).to(torch.int64).to(self._module_device(self.talker))
+        if isinstance(prompt_token_ids_processed, list):
+            prompt_token_ids_processed = (
+                torch.Tensor(prompt_token_ids_processed)
+                .to(torch.int64)
+                .to(self._module_device(self.talker))
+            )
         return prompt_token_ids_processed, prompt_embeds
-            
+
     def _thinker_to_talker_decode_one_step(
         self,
         output_prompt_embeds,
         output_token_ids,
     ):
-        processed_output_token_embeds = output_prompt_embeds + self.talker.get_input_embeddings(output_token_ids) #for decode
+        processed_output_token_embeds = (
+            output_prompt_embeds + self.talker.get_input_embeddings(output_token_ids)
+        )  # for decode
         return output_token_ids, processed_output_token_embeds
 
     def compute_logits(
-        self,
-        hidden_states: Union[torch.Tensor, OmniOutput],
-        sampling_metadata: SamplingMetadata,
+        self, hidden_states: Union[torch.Tensor, OmniOutput]
     ) -> Optional[torch.Tensor]:
         # Handle OmniOutput type
         if isinstance(hidden_states, OmniOutput):
             hidden_states = hidden_states.text_hidden_states
-        
+
         # Use thinker model for logits computation
-        return self.model.compute_logits(hidden_states, sampling_metadata)
+        return self.model.compute_logits(hidden_states)
 
     def sample(
         self,
@@ -511,29 +644,28 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         """
         Generate speech from text tokens using the talker and token2wav models.
         This method is kept for backward compatibility and direct speech generation.
-        
+
         Args:
             text_tokens: Text tokens from thinker model
             voice_type: Voice type for speech generation
-            
+
         Returns:
             Audio tensor
         """
         # Generate codec tokens using talker model
         talker_output = self.talker(
-            input_ids=None,
-            positions=None,
-            inputs_embeds=text_tokens
+            input_ids=None, positions=None, inputs_embeds=text_tokens
         )
-        
+
         # Convert talker output to codec tokens
         codec_tokens = self._convert_to_codec_tokens(talker_output)
-        
+
         # Generate audio using token2wav model
         return self._codec_to_audio(codec_tokens, voice_type=voice_type)
 
-
-    def _convert_to_codec_tokens(self, talker_output: torch.Tensor, sampling_metadata: SamplingMetadata) -> torch.Tensor:
+    def _convert_to_codec_tokens(
+        self, talker_output: torch.Tensor, sampling_metadata: SamplingMetadata
+    ) -> torch.Tensor:
         """
         Reference (HF): use the talker's codec head to obtain logits, suppress BOS,
         then greedily select the next codec token for the current step.
@@ -541,12 +673,19 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         with torch.inference_mode():
             logits = self.talker.compute_logits(talker_output, None)
             if logits is None:
-                return torch.zeros((talker_output.size(0), 0), dtype=torch.long, device=talker_output.device)
+                return torch.zeros(
+                    (talker_output.size(0), 0),
+                    dtype=torch.long,
+                    device=talker_output.device,
+                )
 
-            # Suppress only codec_bos, consistent with HF generate's suppress_tokens behavior
+            # Suppress only codec_bos, consistent with HF generate's
+            # suppress_tokens behavior
             bos_id = None
-            if hasattr(self, 'talker_config') and hasattr(self.talker_config, 'tts_codec_start_token_id'):
-                bos_id = int(getattr(self.talker_config, 'tts_codec_start_token_id'))
+            if hasattr(self, "talker_config") and hasattr(
+                self.talker_config, "tts_codec_start_token_id"
+            ):
+                bos_id = int(getattr(self.talker_config, "tts_codec_start_token_id"))
             if bos_id is not None:
                 logits[..., bos_id] = -1e9
 
@@ -554,21 +693,26 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
             next_id = self.talker.sample(logits, sampling_metadata).sampled_token_ids
             return next_id.to(dtype=torch.long)
 
-    def _init_token2wav_model(self):
-        """Initialize speaker resources if provided; model is constructed in __init__."""
+    def _init_token2wav_model(self, hf_model_folder):
+        """Initialize speaker resources if provided; model is constructed in
+        __init__."""
         if self.token2wav is None or self.token2wav_config is None:
             return
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         # optional speaker resources
-        conds = getattr(self.token2wav_config, 'conds', None)
-        ref_mels = getattr(self.token2wav_config, 'ref_mels', None)
+        conds = getattr(self.token2wav_config, "conds", None)
+        ref_mels = getattr(self.token2wav_config, "ref_mels", None)
         if isinstance(conds, dict) and isinstance(ref_mels, dict):
-            self._token2wav_conds = {k: torch.as_tensor(v, device=device) for k, v in conds.items()}
-            self._token2wav_ref_mels = {k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()}
+            self._token2wav_conds = {
+                k: torch.as_tensor(v, device=device) for k, v in conds.items()
+            }
+            self._token2wav_ref_mels = {
+                k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()
+            }
         # legacy: load from directory if provided
-        model_path = self.vllm_config.model_config.model
+        model_path = hf_model_folder
         if isinstance(model_path, str) and os.path.isdir(model_path):
-            spk_pt = os.path.join(model_path, 'spk_dict.pt')
+            spk_pt = os.path.join(model_path, "spk_dict.pt")
             if os.path.exists(spk_pt):
                 data = torch.load(spk_pt, map_location=device)
                 for key, value in data.items():
@@ -576,20 +720,30 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                     self._token2wav_ref_mels[key] = value["ref_mel"].to(device)
             else:
                 # legacy npy inputs
-                for f in sorted(glob.glob(os.path.join(model_path, 'inputs', '*spk_emb.npy'))):
-                    key = os.path.basename(f).split('_')[0].lower()
-                    self._token2wav_conds[key] = torch.as_tensor(np.load(f), device=device)
-                for f in sorted(glob.glob(os.path.join(model_path, 'inputs', '*ref_mel.npy'))):
-                    key = os.path.basename(f).split('_')[0].lower()
-                    self._token2wav_ref_mels[key] = torch.as_tensor(np.load(f), device=device)
+                for f in sorted(
+                    glob.glob(os.path.join(model_path, "inputs", "*spk_emb.npy"))
+                ):
+                    key = os.path.basename(f).split("_")[0].lower()
+                    self._token2wav_conds[key] = torch.as_tensor(
+                        np.load(f), device=device
+                    )
+                for f in sorted(
+                    glob.glob(os.path.join(model_path, "inputs", "*ref_mel.npy"))
+                ):
+                    key = os.path.basename(f).split("_")[0].lower()
+                    self._token2wav_ref_mels[key] = torch.as_tensor(
+                        np.load(f), device=device
+                    )
 
-    def _codec_to_audio(self, codec_tokens: torch.Tensor, voice_type: str = "default") -> Optional[torch.Tensor]:
+    def _codec_to_audio(
+        self, codec_tokens: torch.Tensor, voice_type: str = "default"
+    ) -> Optional[torch.Tensor]:
         if self.token2wav is None:
             self._init_token2wav_model()
         if self.token2wav is None:
             return None
         # Normalize voice type
-        voice = (voice_type or 'default')
+        voice = voice_type or "default"
         # Resolve cond / ref_mel if provided
         cond = None
         ref_mel = None
@@ -599,9 +753,17 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         # Fallback: create dummy cond/ref_mel if not provided
         token2wav_dev = self._module_device(self.token2wav)
         if cond is None:
-            cond = torch.zeros((1, self.token2wav_config.dit_config.enc_emb_dim), device=token2wav_dev, dtype=torch.float32)
+            cond = torch.zeros(
+                (1, self.token2wav_config.dit_config.enc_emb_dim),
+                device=token2wav_dev,
+                dtype=torch.float32,
+            )
         if ref_mel is None:
-            ref_mel = torch.zeros((1, 300, self.token2wav_config.dit_config.mel_dim), device=token2wav_dev, dtype=torch.float32)
+            ref_mel = torch.zeros(
+                (1, 300, self.token2wav_config.dit_config.mel_dim),
+                device=token2wav_dev,
+                dtype=torch.float32,
+            )
 
         # Ensure codec is (1, T) long tensor on correct device
         if isinstance(codec_tokens, torch.Tensor):
@@ -609,37 +771,49 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
             if codec.ndim == 1:
                 codec = codec.unsqueeze(0)
         else:
-            codec = torch.as_tensor(codec_tokens, dtype=torch.long, device=token2wav_dev).unsqueeze(0)
+            codec = torch.as_tensor(
+                codec_tokens, dtype=torch.long, device=token2wav_dev
+            ).unsqueeze(0)
 
-        # Streaming with chunked process and boundary alignment (rely on token2wav.process_chunk)
-        factor = getattr(self.token2wav.token2wav.factor, 'factor', 2)
+        # Streaming with chunked process and boundary alignment
+        # (rely on token2wav.process_chunk)
+        factor = getattr(self.token2wav.token2wav.factor, "factor", 2)
         chunk_size = 48
-        mel_dim = getattr(self.token2wav.token2wav.code2wav_dit_model, 'mel_dim', self.token2wav_config.dit_config.mel_dim)
+        mel_dim = getattr(
+            self.token2wav.token2wav.code2wav_dit_model,
+            "mel_dim",
+            self.token2wav_config.dit_config.mel_dim,
+        )
         total_mel = int(codec.shape[1] * factor)
         steps = 10
 
         # Prepare initial noise for the whole sequence
-        y_all = torch.randn((1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev)
+        y_all = torch.randn(
+            (1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev
+        )
 
-        logger.info(f"Currently, we do not use the chunked process, we only use the token2wav.process_chunk for the whole sequence.\
-                    The stream mode will be implemented in the future.")
+        logger.info(
+            "Currently, we do not use the chunked process, we only use the "
+            "token2wav.process_chunk for the whole sequence. "
+            "The stream mode will be implemented in the future."
+        )
 
         chunk_ends = []
         for i in range(codec.shape[1]):
             chunk_code_length = i * 2 - 24
-            finished = i==(codec.shape[1]-1)
-            if (chunk_code_length > 0 and
-                chunk_code_length % chunk_size == 0) or finished:
+            finished = i == (codec.shape[1] - 1)
+            if (
+                chunk_code_length > 0 and chunk_code_length % chunk_size == 0
+            ) or finished:
                 chunk_ends.append(i)
 
         # Number of chunks in mel domain
         prev_generated = None
         wav_chunks: list = []
-        prev_id = 0
 
         with torch.inference_mode():
-            for n,i in enumerate([0]):
-                finished = (i == codec.shape[1] - 1)
+            for n, i in enumerate([0]):
+                finished = i == codec.shape[1] - 1
                 _, audio_chunk = self.token2wav.process_chunk(
                     conditioning=cond,
                     reference_mel=ref_mel,
@@ -652,7 +826,6 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 )
                 prev_generated = audio_chunk
                 wav_chunks.append(audio_chunk.detach().cpu().numpy())
-                prev_id = i
 
         if len(wav_chunks) == 0:
             return torch.zeros(0, device=token2wav_dev)
@@ -667,11 +840,11 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         talker_weights = []
         token2wav_weights = []
         for k, v in weights:
-            if k.startswith('thinker.'):
+            if k.startswith("thinker."):
                 thinker_weights.append((k, v))
-            elif k.startswith('talker.'):
+            elif k.startswith("talker."):
                 talker_weights.append((k, v))
-            elif k.startswith('token2wav.'):
+            elif k.startswith("token2wav."):
                 token2wav_weights.append((k, v))
             else:
                 raise ValueError(f"Unknown weight prefix: {k}")
@@ -681,30 +854,52 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
             if thinker_weights:
                 thinker_loaded = self.thinker.load_weights(thinker_weights)
             else:
-                thinker_loaded = set([k for k,v in thinker_weights])
-            thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, 'thinker')
+                thinker_loaded = set([k for k, v in thinker_weights])
+            thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker")
             loaded_weights.update(thinker_loaded)
 
-        
         # Load talker weights
         if talker_weights and self.talker is not None:
             # Map talker weights to appropriate components
             if self.thinker is None:
-                thinker_embedding_weights = [w for n, w in thinker_weights if n == "thinker.model.embed_tokens.weight"]
+                thinker_embedding_weights = [
+                    w
+                    for n, w in thinker_weights
+                    if n == "thinker.model.embed_tokens.weight"
+                ]
                 if thinker_embedding_weights:
-                    self.thinker_embedding = nn.Embedding(thinker_embedding_weights[0].shape[0], thinker_embedding_weights[0].shape[1])
-                    self.thinker_embedding.weight = nn.Parameter(thinker_embedding_weights[0].to(self._module_device(self.talker)))
+                    self.thinker_embedding = nn.Embedding(
+                        thinker_embedding_weights[0].shape[0],
+                        thinker_embedding_weights[0].shape[1],
+                    )
+                    self.thinker_embedding.weight = nn.Parameter(
+                        thinker_embedding_weights[0].to(
+                            self._module_device(self.talker)
+                        )
+                    )
             talker_loaded = self.talker.load_weights(talker_weights)
-            talker_loaded = add_prefix_to_loaded_weights(talker_loaded, 'talker')
+            talker_loaded = add_prefix_to_loaded_weights(talker_loaded, "talker")
             loaded_weights.update(talker_loaded)
             loaded_weights.update(self._init_special_tokens_embeddings())
-            
-        
+
         # Load token2wav weights (if any)
         if token2wav_weights and self.token2wav is not None:
-            self._init_token2wav_model()
-            t2w_loaded = self.token2wav.load_weights(token2wav_weights, os.path.join(self.vllm_config.model_config.model, "spk_dict.pt"))
-            t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, 'token2wav')
+            # download weights from huggingface for spk_dict.pt
+            model_path = self.vllm_config.model_config.model
+            download_dir = self.vllm_config.load_config.download_dir
+            if os.path.exists(model_path):
+                hf_model_folder = model_path
+            else:
+                hf_model_folder = download_weights_from_hf_specific(
+                    model_path,
+                    download_dir,
+                    allow_patterns=["*.pt"],
+                )
+            self._init_token2wav_model(hf_model_folder)
+            t2w_loaded = self.token2wav.load_weights(
+                token2wav_weights, os.path.join(hf_model_folder, "spk_dict.pt")
+            )
+            t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, "token2wav")
             loaded_weights.update(t2w_loaded)
-        
-        return loaded_weights
+
+        return loaded_weights

vllm_omni/model_executor/models/qwen2_5_omni_talker.py

@@ -1,32 +1,43 @@
 from functools import cached_property
-from typing import Iterable, List, Optional, Set, Tuple, Union
+from typing import Iterable, Optional, Set, Tuple, Union
 
 import torch
 import torch.nn as nn
-from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniTalkerConfig
-from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
+from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
+    Qwen2_5OmniTalkerConfig,
+)
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
+    Qwen2_5OmniAudioEncoder,
+)
 
-from vllm.attention import AttentionMetadata
+# from vllm.attention import AttentionMetadata  # unused import
 from vllm.config import VllmConfig
-from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ColumnParallelLinear
-from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
-from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+)
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
     Qwen2_5OmniThinkerProcessingInfo,
-    Qwen2_5OmniThinkerDummyInputsBuilder)
-from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.logger import init_logger
-from vllm.sequence import IntermediateTensors
-
-from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP, MultiModalEmbeddings
-from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
-                    init_vllm_registered_model,
-                    maybe_prefix,
-                    merge_multimodal_embeddings)
+)
 from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    init_vllm_registered_model,
+    maybe_prefix,
+    merge_multimodal_embeddings,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
 
 
 @MULTIMODAL_REGISTRY.register_processor(
@@ -34,9 +45,9 @@ from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer
     info=Qwen2_5OmniThinkerProcessingInfo,
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
-class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
-                                                SupportsPP,
-                                                Qwen2_5OmniConditionalGenerationMixin):
+class Qwen2_5OmniTalkerForConditionalGeneration(
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+):
     logger = init_logger(__name__)
     # Align to thinker-style static mapper for clarity
     hf_to_vllm_mapper = WeightsMapper(
@@ -48,7 +59,8 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
             "talker.thinker_to_talker_proj.": "thinker_to_talker_proj.",
             # fallback root
             "talker.": "",
-        })
+        }
+    )
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -60,8 +72,9 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
 
         if hasattr(config, "talker_config"):
             self.config = config.talker_config
-            vllm_config.model_config.hf_text_config = \
-                    vllm_config.model_config.hf_config.talker_config
+            vllm_config.model_config.hf_text_config = (
+                vllm_config.model_config.hf_config.talker_config
+            )
         else:
             self.config = config
 
@@ -76,12 +89,13 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
         self.language_model = init_vllm_registered_model(
             vllm_config=vllm_config,
             prefix=maybe_prefix(prefix, "language_model"),
-            hf_config=getattr(self.config, 'text_config', self.config),
+            hf_config=getattr(self.config, "text_config", self.config),
             architectures=["Qwen2ForCausalLM_old"],
         )
         self.make_empty_intermediate_tensors = (
-            self.language_model.make_empty_intermediate_tensors)
-    
+            self.language_model.make_empty_intermediate_tensors
+        )
+
     def init_multi_modal(self, thinker_config):
         self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
         self.visual = Qwen2_5_VisionTransformer(
@@ -99,7 +113,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
         if hasattr(self.language_model, "sampler"):
             return self.language_model.sampler
 
-        return get_sampler()
+        return Sampler()
 
     def get_input_embeddings(
         self,
@@ -107,17 +121,20 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
         multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.get_input_embeddings(input_ids)
-        if multimodal_embeddings is not None \
-            and len(multimodal_embeddings) != 0:
+        if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
 
             # TODO (ywang96): support overlapping modalitiy embeddings so that
             # `use_audio_in_video` will work on V1.
             inputs_embeds = merge_multimodal_embeddings(
-                input_ids, inputs_embeds, multimodal_embeddings, [
+                input_ids,
+                inputs_embeds,
+                multimodal_embeddings,
+                [
                     self.config.image_token_index,
                     self.config.video_token_index,
-                    self.config.audio_token_index
-                ])
+                    self.config.audio_token_index,
+                ],
+            )
         return inputs_embeds
 
     def forward(
@@ -128,9 +145,10 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
         inputs_embeds: Optional[torch.Tensor] = None,
         **kwargs: object,
     ) -> Union[torch.Tensor, IntermediateTensors]:
-        assert input_ids is not None or inputs_embeds is not None, "input_ids or inputs_embeds must be provided"
-        forward_context: ForwardContext = get_forward_context()
-        attn_metadata: AttentionMetadata = forward_context.attn_metadata
+        assert (
+            input_ids is not None or inputs_embeds is not None
+        ), "input_ids or inputs_embeds must be provided"
+        # forward_context: ForwardContext = get_forward_context()  # unused variable
 
         if intermediate_tensors is not None:
             inputs_embeds = None
@@ -143,19 +161,13 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
         # projection
         inputs_embeds, _ = self.thinker_to_talker_proj(inputs_embeds)
 
-        hidden_states = self.language_model.model(input_ids,
-                                                  positions,
-                                                  intermediate_tensors,
-                                                  inputs_embeds=inputs_embeds)
+        hidden_states = self.language_model.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
+        )
         return hidden_states
 
-    def compute_logits(
-        self,
-        hidden_states: torch.Tensor,
-        sampling_metadata: SamplingMetadata,
-    ) -> Optional[torch.Tensor]:
-        return self.language_model.compute_logits(hidden_states,
-                                                  sampling_metadata)
+    def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:
+        return self.language_model.compute_logits(hidden_states)
 
     def sample(
         self,
@@ -164,11 +176,10 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
     ) -> Optional[SamplerOutput]:
         return self.language_model.sample(logits, sampling_metadata)
 
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         loader = AutoWeightsLoader(
             self,
-            skip_prefixes=['thinker.', 'token2wav.'],
+            skip_prefixes=["thinker.", "token2wav."],
         )
         loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
         # Log load summary
@@ -180,42 +191,53 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
             device = next(self.parameters()).device
             self.logger.info(
                 "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
-                self.__class__.__name__, True, total_bytes / (1024**2), str(device))
+                self.__class__.__name__,
+                True,
+                total_bytes / (1024**2),
+                str(device),
+            )
         except Exception:
             pass
         multi_model_weights = set()
         for name, param in self.visual.named_parameters():
-            multi_model_weights.add("visual."+name)
+            multi_model_weights.add("visual." + name)
         for name, param in self.audio_tower.named_parameters():
-            multi_model_weights.add("audio_tower."+name)
+            multi_model_weights.add("audio_tower." + name)
         loaded.update(multi_model_weights)
         return loaded
-    
+
     def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
         mm_input_by_modality = {}
 
         # Preserve the order of modalities if there are multiple of them
         # from the order of kwargs.
         for input_key in kwargs:
-            if input_key in ("pixel_values", "image_embeds"
-                             ) and "image" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "image"] = self._parse_and_validate_image_input(**kwargs)
-            if input_key in ("pixel_values_videos", "video_embeds"
-                             ) and "video" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "video"] = self._parse_and_validate_video_input(**kwargs)
-            if input_key in ("input_audio_features"
-                             ) and "audio" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "audio"] = self._parse_and_validate_audio_input(**kwargs)
+            if (
+                input_key in ("pixel_values", "image_embeds")
+                and "image" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("pixel_values_videos", "video_embeds")
+                and "video" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["video"] = self._parse_and_validate_video_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("input_audio_features")
+                and "audio" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
+                    **kwargs
+                )
         return mm_input_by_modality
 
-    def get_multimodal_embeddings(self,
-                                  **kwargs: object) -> MultiModalEmbeddings:
+    def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
 
-        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(
-            **kwargs)
+        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []
 
@@ -236,4 +258,4 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
             if modality == "audio":
                 audio_embeddings = self._process_audio_input(multimodal_input)
                 multimodal_embeddings += audio_embeddings
-        return multimodal_embeddings
+        return multimodal_embeddings

vllm_omni/model_executor/models/qwen2_5_omni_thinker.py

@@ -7,45 +7,72 @@ import torch
 import torch.nn as nn
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
-    Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig)
+    Qwen2_5OmniConfig,
+    Qwen2_5OmniThinkerConfig,
+)
 from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
-    Qwen2_5OmniAudioEncoder)
+    Qwen2_5OmniAudioEncoder,
+)
 from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import (
-    Qwen2_5OmniProcessor)
+    Qwen2_5OmniProcessor,
+)
 from transformers.models.whisper import WhisperFeatureExtractor
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+)
 from vllm.model_executor.models.qwen2_5_vl import (
-    Qwen2_5_VisionTransformer, Qwen2_5_VLImageEmbeddingInputs,
-    Qwen2_5_VLImageInputs, Qwen2_5_VLImagePixelInputs,
-    Qwen2_5_VLProcessingInfo, Qwen2_5_VLVideoEmbeddingInputs,
-    Qwen2_5_VLVideoInputs, Qwen2_5_VLVideoPixelInputs)
+    Qwen2_5_VisionTransformer,
+    Qwen2_5_VLImageEmbeddingInputs,
+    Qwen2_5_VLImageInputs,
+    Qwen2_5_VLImagePixelInputs,
+    Qwen2_5_VLProcessingInfo,
+    Qwen2_5_VLVideoEmbeddingInputs,
+    Qwen2_5_VLVideoInputs,
+    Qwen2_5_VLVideoPixelInputs,
+)
 from vllm.model_executor.models.qwen2_audio import (
-    Qwen2AudioInputs, Qwen2AudioProcessingInfo,
-    _get_feat_extract_output_lengths)
+    Qwen2AudioInputs,
+    Qwen2AudioProcessingInfo,
+    _get_feat_extract_output_lengths,
+)
 from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser
-from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    init_vllm_registered_model,
+    maybe_prefix,
+    merge_multimodal_embeddings,
+)
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import (ImageItem, ModalityData,
-                                    MultiModalDataDict, MultiModalFieldConfig,
-                                    MultiModalKwargs, NestedTensors)
-from vllm.multimodal.parse import (AudioProcessorItems, DictEmbeddingItems,
-                                   ModalityDataItems, MultiModalDataItems,
-                                   MultiModalDataParser)
-from vllm.multimodal.processing import (BaseMultiModalProcessor,
-                                        PlaceholderFeaturesInfo,
-                                        PromptReplacement, PromptUpdate)
+from vllm.multimodal.inputs import (
+    ImageItem,
+    ModalityData,
+    MultiModalDataDict,
+    MultiModalFieldConfig,
+    MultiModalKwargs,
+    NestedTensors,
+)
+from vllm.multimodal.parse import (
+    AudioProcessorItems,
+    DictEmbeddingItems,
+    ModalityDataItems,
+    MultiModalDataItems,
+    MultiModalDataParser,
+)
+from vllm.multimodal.processing import (
+    BaseMultiModalProcessor,
+    PlaceholderFeaturesInfo,
+    PromptReplacement,
+    PromptUpdate,
+)
 from vllm.multimodal.profiling import BaseDummyInputsBuilder
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.tokenizer import decode_tokens, encode_tokens
-
-from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
-from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
-                    init_vllm_registered_model, maybe_prefix,
-                    merge_multimodal_embeddings)
-
 from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
 
 try:
@@ -57,8 +84,7 @@ logger = init_logger(__name__)
 
 
 def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
-    audio_feature_lengths = hf_inputs.get("audio_feature_lengths",
-                                          torch.empty((0, )))
+    audio_feature_lengths = hf_inputs.get("audio_feature_lengths", torch.empty((0,)))
 
     image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
     image_grid_sizes = image_grid_thw.prod(-1)
@@ -68,18 +94,17 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
 
     return dict(
         input_audio_features=MultiModalFieldConfig.flat_from_sizes(
-            "audio", audio_feature_lengths, dim=1),
+            "audio", audio_feature_lengths, dim=1
+        ),
         feature_attention_mask=MultiModalFieldConfig.batched("audio"),
         audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
-        pixel_values=MultiModalFieldConfig.flat_from_sizes(
-            "image", image_grid_sizes),
-        image_embeds=MultiModalFieldConfig.flat_from_sizes(
-            "image", image_grid_sizes),
+        pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
+        image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
         image_grid_thw=MultiModalFieldConfig.batched("image"),
         pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
-            "video", video_grid_sizes),
-        video_embeds=MultiModalFieldConfig.flat_from_sizes(
-            "video", video_grid_sizes),
+            "video", video_grid_sizes
+        ),
+        video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes),
         video_grid_thw=MultiModalFieldConfig.batched("video"),
         second_per_grid_ts=MultiModalFieldConfig.batched("video"),
     )
@@ -95,17 +120,16 @@ class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
             return DictEmbeddingItems(
                 data,
                 modality="audio",
-                required_fields={
-                    "input_audio_features", "audio_feature_lengths"
-                },
+                required_fields={"input_audio_features", "audio_feature_lengths"},
                 fields_factory=_qwen2_5_omni_thinker_field_config,
             )
 
         return super()._parse_audio_data(data)
 
 
-class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo,
-                                       Qwen2_5_VLProcessingInfo):
+class Qwen2_5OmniThinkerProcessingInfo(
+    Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo
+):
 
     def get_hf_config(self):
         return self.ctx.get_hf_config(Qwen2_5OmniConfig).thinker_config
@@ -128,7 +152,8 @@ class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo,
 
 
 class Qwen2_5OmniThinkerDummyInputsBuilder(
-        BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]):
+    BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]
+):
 
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
         num_audios = mm_counts.get("audio", 0)
@@ -141,8 +166,11 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
         image_token: str = hf_processor.image_token
         video_token: str = hf_processor.video_token
 
-        return (audio_token * num_audios + image_token * num_images +
-                video_token * num_videos)
+        return (
+            audio_token * num_audios
+            + image_token * num_images
+            + video_token * num_videos
+        )
 
     def get_dummy_mm_data(
         self,
@@ -155,40 +183,45 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
 
         feature_extractor = self.info.get_feature_extractor()
 
-        target_audio_length = min(
-            feature_extractor.chunk_length,
-            30,
-        ) * feature_extractor.sampling_rate
-        target_width, target_height = \
-            self.info.get_image_size_with_most_features()
-        target_num_frames = \
-            self.info.get_num_frames_with_most_features(seq_len, mm_counts)
+        target_audio_length = (
+            min(
+                feature_extractor.chunk_length,
+                30,
+            )
+            * feature_extractor.sampling_rate
+        )
+        target_width, target_height = self.info.get_image_size_with_most_features()
+        target_num_frames = self.info.get_num_frames_with_most_features(
+            seq_len, mm_counts
+        )
 
         mm_data = {
-            "audio":
-            self._get_dummy_audios(length=target_audio_length,
-                                   num_audios=num_audios),
-            "image":
-            self._get_dummy_images(width=target_width,
-                                   height=target_height,
-                                   num_images=num_images),
-            "video":
-            self._get_dummy_videos(width=target_width,
-                                   height=target_height,
-                                   num_frames=target_num_frames,
-                                   num_videos=num_videos),
+            "audio": self._get_dummy_audios(
+                length=target_audio_length, num_audios=num_audios
+            ),
+            "image": self._get_dummy_images(
+                width=target_width, height=target_height, num_images=num_images
+            ),
+            "video": self._get_dummy_videos(
+                width=target_width,
+                height=target_height,
+                num_frames=target_num_frames,
+                num_videos=num_videos,
+            ),
         }
 
         return mm_data
 
 
 class Qwen2_5OmniThinkerMultiModalProcessor(
-        BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]):
+    BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]
+):
 
     def _get_data_parser(self) -> MultiModalDataParser:
         feature_extractor = self.info.get_feature_extractor()
         return Qwen2_5OmniThinkerMultiModalDataParser(
-            target_sr=feature_extractor.sampling_rate)
+            target_sr=feature_extractor.sampling_rate
+        )
 
     def _call_hf_processor(
         self,
@@ -204,7 +237,9 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         if audios:
             # NOTE: Qwen2.5-Omni processor accept "audio"
             mm_data["audio"] = audios
-            mm_kwargs = dict(**mm_kwargs, )
+            mm_kwargs = dict(
+                **mm_kwargs,
+            )
 
         hf_inputs = super()._call_hf_processor(
             prompt=prompt,
@@ -213,17 +248,19 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             tok_kwargs=tok_kwargs,
         )
 
-        input_features = hf_inputs.pop('input_features', None)
-        feature_attention_mask = hf_inputs.get('feature_attention_mask', None)
-        if ('input_audio_features' not in hf_inputs
-                and input_features is not None):
+        input_features = hf_inputs.pop("input_features", None)
+        feature_attention_mask = hf_inputs.get("feature_attention_mask", None)
+        if "input_audio_features" not in hf_inputs and input_features is not None:
             if feature_attention_mask is not None:
-                input_features = input_features.permute(
-                    0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
-            hf_inputs['input_audio_features'] = input_features
-        if ('audio_feature_lengths' not in hf_inputs
-                and feature_attention_mask is not None):
-            hf_inputs['audio_feature_lengths'] = feature_attention_mask.sum(-1)
+                input_features = input_features.permute(0, 2, 1)[
+                    feature_attention_mask.bool()
+                ].permute(1, 0)
+            hf_inputs["input_audio_features"] = input_features
+        if (
+            "audio_feature_lengths" not in hf_inputs
+            and feature_attention_mask is not None
+        ):
+            hf_inputs["audio_feature_lengths"] = feature_attention_mask.sum(-1)
         return hf_inputs
 
     def _get_mm_fields_config(
@@ -249,14 +286,12 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             hf_processor_mm_kwargs,
             mm_kwargs,
         )
-        mm_prompt_updates = self._bind_and_group_updates(
-            unbound_prompt_updates)
+        mm_prompt_updates = self._bind_and_group_updates(unbound_prompt_updates)
 
         mm_item_counts = mm_items.get_all_counts()
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
 
-        use_audio_in_video = hf_processor_mm_kwargs.get(
-            "use_audio_in_video", False)
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
 
         if is_update_applied:
             mm_placeholders = self._find_mm_placeholders(
@@ -265,9 +300,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 mm_item_counts,
             )
             self._validate_mm_placeholders(
-                mm_placeholders,
-                mm_item_counts,
-                use_audio_in_video=use_audio_in_video)
+                mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
+            )
 
             tokenizer = self.info.get_tokenizer()
             prompt = decode_tokens(tokenizer, prompt_ids)
@@ -282,9 +316,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 mm_item_counts,
             )
             self._validate_mm_placeholders(
-                mm_placeholders,
-                mm_item_counts,
-                use_audio_in_video=use_audio_in_video)
+                mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
+            )
 
         tokenizer = self.info.get_tokenizer()
         prompt = decode_tokens(tokenizer, prompt_ids)
@@ -302,8 +335,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
     ) -> Sequence[PromptUpdate]:
         processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
         tokenizer = self.info.get_tokenizer()
-        image_processor = self.info.get_image_processor(
-            **hf_processor_mm_kwargs)
+        image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)
         vocab = tokenizer.get_vocab()
 
         audio_token = processor.audio_token
@@ -319,12 +351,14 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             audio_output_lengths = []
         elif audio_feature_lengths is not None:
             _, audio_output_lens = _get_feat_extract_output_lengths(
-                audio_feature_lengths)
+                audio_feature_lengths
+            )
             audio_output_lengths = audio_output_lens.tolist()
         elif feature_attention_mask is not None:
             assert isinstance(feature_attention_mask, torch.Tensor)
             _, audio_output_lens = _get_feat_extract_output_lengths(
-                feature_attention_mask.sum(-1))
+                feature_attention_mask.sum(-1)
+            )
             audio_output_lengths = audio_output_lens.tolist()
 
         # number of audios read from video.
@@ -339,7 +373,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 audio = audios.get(item_idx)
                 raise ValueError(
                     f"The audio {audio} (len={len(audio)}) is too short "
-                    "to be represented inside the model")
+                    "to be represented inside the model"
+                )
 
             return [audio_token_id] * num_features
 
@@ -351,21 +386,20 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             token_id = image_token_id if modality == "image" else video_token_id
             return [token_id] * (int(grid_thw.prod()) // merge_length)
 
-        use_audio_in_video = hf_processor_mm_kwargs.get(
-            "use_audio_in_video", False)
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
         thinker_config = self.info.get_hf_config()
 
         def get_replacement_qwen2_use_audio_in_video(item_idx: int):
             nonlocal audio_in_video_item_idx
 
-            audio_num_features = audio_output_lengths[audio_in_video_item_idx +
-                                                      item_idx]
+            audio_num_features = audio_output_lengths[
+                audio_in_video_item_idx + item_idx
+            ]
             video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx]
 
             audio_in_video_item_idx += 1
 
-            second_per_grid_ts = hf_processor_mm_kwargs.get(
-                "second_per_grid_ts", None)
+            second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
             if second_per_grid_ts:
                 video_second_per_grid_t = second_per_grid_ts[item_idx]
             else:
@@ -379,8 +413,10 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             )
 
         video_replacement_fn = (
-            get_replacement_qwen2_use_audio_in_video if use_audio_in_video else
-            partial(get_replacement_qwen2_vision, modality="video"))
+            get_replacement_qwen2_use_audio_in_video
+            if use_audio_in_video
+            else partial(get_replacement_qwen2_vision, modality="video")
+        )
 
         return [
             PromptReplacement(
@@ -391,8 +427,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             PromptReplacement(
                 modality="image",
                 target=image_token,
-                replacement=partial(get_replacement_qwen2_vision,
-                                    modality="image"),
+                replacement=partial(get_replacement_qwen2_vision, modality="image"),
             ),
             PromptReplacement(
                 modality="video",
@@ -445,8 +480,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         """
         mm_counts = mm_items.get_all_counts()
 
-        use_audio_in_video = hf_processor_mm_kwargs.get(
-            "use_audio_in_video", False)
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
         if use_audio_in_video and "video" in mm_counts:
             assert "audio" in mm_counts
             mm_counts["audio"] -= mm_counts["video"]
@@ -476,36 +510,43 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
 
 class Qwen2_5OmniConditionalGenerationMixin:
 
-    def _validate_and_reshape_mm_tensor(self,
-                                        mm_input: object,
-                                        name: str,
-                                        dim: int = 0) -> torch.Tensor:
+    def _validate_and_reshape_mm_tensor(
+        self, mm_input: object, name: str, dim: int = 0
+    ) -> torch.Tensor:
         if not isinstance(mm_input, (torch.Tensor, list)):
-            raise ValueError(f"Incorrect type of {name}. "
-                             f"Got type: {type(mm_input)}")
+            raise ValueError(
+                f"Incorrect type of {name}. " f"Got type: {type(mm_input)}"
+            )
         if isinstance(mm_input, torch.Tensor):
             return torch.concat(list(mm_input), dim=dim)
         else:
             return torch.concat(mm_input, dim=dim)
 
     def _parse_and_validate_audio_input(
-            self, **kwargs: object) -> Optional[Qwen2AudioInputs]:
-        input_audio_features = kwargs.pop('input_audio_features', None)
-        audio_feature_lengths = kwargs.pop('audio_feature_lengths', None)
-        feature_attention_mask = kwargs.pop('feature_attention_mask', None)
+        self, **kwargs: object
+    ) -> Optional[Qwen2AudioInputs]:
+        input_audio_features = kwargs.pop("input_audio_features", None)
+        audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
+        feature_attention_mask = kwargs.pop("feature_attention_mask", None)
         if input_audio_features is None:
             return None
         input_audio_features = self._validate_and_reshape_mm_tensor(
-            input_audio_features, 'input_audio_features', dim=1)
+            input_audio_features, "input_audio_features", dim=1
+        )
         if feature_attention_mask is not None:
             feature_attention_mask = self._validate_and_reshape_mm_tensor(
-                feature_attention_mask, 'feature_attention_mask')
+                feature_attention_mask, "feature_attention_mask"
+            )
         if not isinstance(input_audio_features, (torch.Tensor, list)):
-            raise ValueError("Incorrect type of audio input features. "
-                             f"Got type: {type(input_audio_features)}")
-        return Qwen2AudioInputs(input_features=input_audio_features,
-                                audio_feature_lengths=audio_feature_lengths,
-                                feature_attention_mask=feature_attention_mask)
+            raise ValueError(
+                "Incorrect type of audio input features. "
+                f"Got type: {type(input_audio_features)}"
+            )
+        return Qwen2AudioInputs(
+            input_features=input_audio_features,
+            audio_feature_lengths=audio_feature_lengths,
+            feature_attention_mask=feature_attention_mask,
+        )
 
     def _parse_and_validate_image_input(
         self,
@@ -520,31 +561,42 @@ class Qwen2_5OmniConditionalGenerationMixin:
 
         if pixel_values is not None:
             pixel_values = self._validate_and_reshape_mm_tensor(
-                pixel_values, "image pixel values")
+                pixel_values, "image pixel values"
+            )
             image_grid_thw = self._validate_and_reshape_mm_tensor(
-                image_grid_thw, "image grid_thw")
+                image_grid_thw, "image grid_thw"
+            )
 
             if not isinstance(pixel_values, (torch.Tensor, list)):
-                raise ValueError("Incorrect type of image pixel values. "
-                                 f"Got type: {type(pixel_values)}")
+                raise ValueError(
+                    "Incorrect type of image pixel values. "
+                    f"Got type: {type(pixel_values)}"
+                )
 
-            return Qwen2_5_VLImagePixelInputs(type="pixel_values",
-                                              pixel_values=pixel_values,
-                                              image_grid_thw=image_grid_thw)
+            return Qwen2_5_VLImagePixelInputs(
+                type="pixel_values",
+                pixel_values=pixel_values,
+                image_grid_thw=image_grid_thw,
+            )
 
         if image_embeds is not None:
             image_embeds = self._validate_and_reshape_mm_tensor(
-                image_embeds, "image embeds")
+                image_embeds, "image embeds"
+            )
             image_grid_thw = self._validate_and_reshape_mm_tensor(
-                image_grid_thw, "image grid_thw")
+                image_grid_thw, "image grid_thw"
+            )
 
             if not isinstance(image_embeds, torch.Tensor):
-                raise ValueError("Incorrect type of image embeddings. "
-                                 f"Got type: {type(image_embeds)}")
+                raise ValueError(
+                    "Incorrect type of image embeddings. "
+                    f"Got type: {type(image_embeds)}"
+                )
             return Qwen2_5_VLImageEmbeddingInputs(
                 type="image_embeds",
                 image_embeds=image_embeds,
-                image_grid_thw=image_grid_thw)
+                image_grid_thw=image_grid_thw,
+            )
 
     def _parse_and_validate_video_input(
         self,
@@ -559,9 +611,11 @@ class Qwen2_5OmniConditionalGenerationMixin:
 
         if pixel_values_videos is not None:
             pixel_values_videos = self._validate_and_reshape_mm_tensor(
-                pixel_values_videos, "video pixel values")
+                pixel_values_videos, "video pixel values"
+            )
             video_grid_thw = self._validate_and_reshape_mm_tensor(
-                video_grid_thw, "video grid_thw")
+                video_grid_thw, "video grid_thw"
+            )
 
             return Qwen2_5_VLVideoPixelInputs(
                 type="pixel_values_videos",
@@ -571,17 +625,22 @@ class Qwen2_5OmniConditionalGenerationMixin:
 
         if video_embeds is not None:
             video_embeds = self._validate_and_reshape_mm_tensor(
-                video_embeds, "video embeds")
+                video_embeds, "video embeds"
+            )
             video_grid_thw = self._validate_and_reshape_mm_tensor(
-                video_grid_thw, "video grid_thw")
+                video_grid_thw, "video grid_thw"
+            )
 
             if not isinstance(video_embeds, torch.Tensor):
-                raise ValueError("Incorrect type of video embeddings. "
-                                 f"Got type: {type(video_embeds)}")
+                raise ValueError(
+                    "Incorrect type of video embeddings. "
+                    f"Got type: {type(video_embeds)}"
+                )
             return Qwen2_5_VLVideoEmbeddingInputs(
                 type="video_embeds",
                 video_embeds=video_embeds,
-                video_grid_thw=video_grid_thw)
+                video_grid_thw=video_grid_thw,
+            )
 
     def _process_audio_input(
         self,
@@ -596,16 +655,18 @@ class Qwen2_5OmniConditionalGenerationMixin:
             assert input_features.shape[0] == 1
             input_features = input_features.squeeze(0)
         if audio_feature_lengths.ndim == 2:
-            assert audio_feature_lengths.shape[
-                0] == 1 or audio_feature_lengths.shape[1] == 1
+            assert (
+                audio_feature_lengths.shape[0] == 1
+                or audio_feature_lengths.shape[1] == 1
+            )
             if audio_feature_lengths.shape[0] == 1:
                 audio_feature_lengths = audio_feature_lengths.squeeze(0)
             else:
                 audio_feature_lengths = audio_feature_lengths.squeeze(1)
 
         audio_feat_lengths, audio_output_lengths = (
-            self.audio_tower._get_feat_extract_output_lengths(
-                audio_feature_lengths))
+            self.audio_tower._get_feat_extract_output_lengths(audio_feature_lengths)
+        )
 
         audio_outputs = self.audio_tower(
             input_features.to(self.audio_tower.dtype),
@@ -616,8 +677,8 @@ class Qwen2_5OmniConditionalGenerationMixin:
         return audio_features.split(audio_output_lengths.tolist())
 
     def _process_image_input(
-            self,
-            image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]:
+        self, image_input: Qwen2_5_VLImageInputs
+    ) -> tuple[torch.Tensor, ...]:
         if image_input["type"] == "image_embeds":
             return image_input["image_embeds"].type(self.visual.dtype)
 
@@ -633,18 +694,18 @@ class Qwen2_5OmniConditionalGenerationMixin:
         return image_embeds.split(sizes.tolist())
 
     def _process_video_input(
-            self,
-            video_input: Qwen2_5_VLVideoInputs,
-            video_hashes: list[str] = None,
-            cached_video_embeds: torch.Tensor = None) -> torch.Tensor:
+        self,
+        video_input: Qwen2_5_VLVideoInputs,
+        video_hashes: list[str] = None,
+        cached_video_embeds: torch.Tensor = None,
+    ) -> torch.Tensor:
         if video_input["type"] == "video_embeds":
             return video_input["video_embeds"].type(self.visual.dtype)
 
         grid_thw = video_input["video_grid_thw"]
         assert grid_thw.ndim == 2
 
-        pixel_values_videos = video_input["pixel_values_videos"].type(
-            self.visual.dtype)
+        pixel_values_videos = video_input["pixel_values_videos"].type(self.visual.dtype)
         video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)
         # Split concatenated embeddings for each video item.
         merge_size = self.visual.spatial_merge_size
@@ -659,14 +720,15 @@ class Qwen2_5OmniConditionalGenerationMixin:
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniThinkerForConditionalGeneration(
-        nn.Module, SupportsMultiModal, SupportsPP,
-        Qwen2_5OmniConditionalGenerationMixin):
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+):
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "thinker.lm_head.": "language_model.lm_head.",
             "thinker.model.": "language_model.model.",
             "thinker.": "",
-        })
+        }
+    )
 
     @classmethod
     def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
@@ -681,8 +743,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        thinker_config: Qwen2_5OmniThinkerConfig = (
-            vllm_config.model_config.hf_config)
+        thinker_config: Qwen2_5OmniThinkerConfig = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
         multimodal_config = vllm_config.model_config.multimodal_config
         self.config = thinker_config
@@ -698,7 +759,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             logger.warning(
                 "flash_attn is not available, the model may not yield the "
                 "exactly same result as the transformers implementation "
-                "in the audio tower part.")
+                "in the audio tower part."
+            )
 
         self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
         self.visual = Qwen2_5_VisionTransformer(
@@ -716,7 +778,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         )
 
         self.make_empty_intermediate_tensors = (
-            self.language_model.make_empty_intermediate_tensors)
+            self.language_model.make_empty_intermediate_tensors
+        )
 
     def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
         mm_input_by_modality = {}
@@ -724,28 +787,35 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         # Preserve the order of modalities if there are multiple of them
         # from the order of kwargs.
         for input_key in kwargs:
-            if input_key in ("pixel_values", "image_embeds"
-                             ) and "image" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "image"] = self._parse_and_validate_image_input(**kwargs)
-            if input_key in ("pixel_values_videos", "video_embeds"
-                             ) and "video" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "video"] = self._parse_and_validate_video_input(**kwargs)
-            if input_key in ("input_audio_features"
-                             ) and "audio" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "audio"] = self._parse_and_validate_audio_input(**kwargs)
+            if (
+                input_key in ("pixel_values", "image_embeds")
+                and "image" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("pixel_values_videos", "video_embeds")
+                and "video" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["video"] = self._parse_and_validate_video_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("input_audio_features")
+                and "audio" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
+                    **kwargs
+                )
         return mm_input_by_modality
 
     def get_language_model(self) -> torch.nn.Module:
         return self.language_model
 
-    def get_multimodal_embeddings(self,
-                                  **kwargs: object) -> MultiModalEmbeddings:
+    def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
 
-        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(
-            **kwargs)
+        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []
 
@@ -774,21 +844,23 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.get_input_embeddings(input_ids)
-        if multimodal_embeddings is not None \
-            and len(multimodal_embeddings) != 0:
+        if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
 
             # TODO (ywang96): support overlapping modalitiy embeddings so that
             # `use_audio_in_video` will work on V1.
             inputs_embeds = merge_multimodal_embeddings(
-                input_ids, inputs_embeds, multimodal_embeddings, [
+                input_ids,
+                inputs_embeds,
+                multimodal_embeddings,
+                [
                     self.config.image_token_index,
                     self.config.video_token_index,
-                    self.config.audio_token_index
-                ])
+                    self.config.audio_token_index,
+                ],
+            )
         return inputs_embeds
 
-    def get_multimodal_embeddings_v0(
-            self, **kwargs: object) -> Optional[NestedTensors]:
+    def get_multimodal_embeddings_v0(self, **kwargs: object) -> Optional[NestedTensors]:
         audio_input = self._parse_and_validate_audio_input(**kwargs)
         image_input = self._parse_and_validate_image_input(**kwargs)
         video_input = self._parse_and_validate_video_input(**kwargs)
@@ -826,7 +898,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             if modality == "video":
                 placeholder_token_id = self.config.video_token_index
             inputs_embeds = merge_multimodal_embeddings(
-                input_ids, inputs_embeds, embeddings, placeholder_token_id)
+                input_ids, inputs_embeds, embeddings, placeholder_token_id
+            )
         return inputs_embeds
 
     def forward(
@@ -844,35 +917,36 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         # condition is for v0 compatibility.
         elif inputs_embeds is None:
             multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)
-            inputs_embeds = self.get_input_embeddings(
-                input_ids, multimodal_embeddings)
-            text_inputs_embeds = self.get_input_embeddings(input_ids, [(torch.zeros_like(embeddings), modality) for embeddings, modality in multimodal_embeddings] if multimodal_embeddings is not None else None)
+            inputs_embeds = self.get_input_embeddings(input_ids, multimodal_embeddings)
+            text_inputs_embeds = self.get_input_embeddings(
+                input_ids,
+                (
+                    [
+                        (torch.zeros_like(embeddings), modality)
+                        for embeddings, modality in multimodal_embeddings
+                    ]
+                    if multimodal_embeddings is not None
+                    else None
+                ),
+            )
             input_ids = None
         else:
             text_inputs_embeds = inputs_embeds
 
-        hidden_states = self.language_model.model(input_ids,
-                                                  positions,
-                                                  intermediate_tensors,
-                                                  inputs_embeds=inputs_embeds)
-        return text_inputs_embeds, hidden_states.unsqueeze(0) #(1, S, D)
+        hidden_states = self.language_model.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
+        )
+        return text_inputs_embeds, hidden_states.unsqueeze(0)  # (1, S, D)
 
-    def compute_logits(
-        self,
-        hidden_states: torch.Tensor,
-        sampling_metadata: SamplingMetadata,
-    ) -> Optional[torch.Tensor]:
-        return self.language_model.compute_logits(hidden_states,
-                                                  sampling_metadata)
+    def compute_logits(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]:
+        return self.language_model.compute_logits(hidden_states)
 
-    def load_weights(self, weights: Iterable[tuple[str,
-                                                   torch.Tensor]]) -> set[str]:
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(
             self,
             skip_prefixes=["talker.", "token2wav."],
         )
-        loaded_weights = loader.load_weights(weights,
-                                             mapper=self.hf_to_vllm_mapper)
+        loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
         # Log load summary
         try:
             total_bytes = 0
@@ -882,8 +956,11 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             device = next(self.parameters()).device
             logger.info(
                 "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
-                self.__class__.__name__, True, total_bytes / (1024**2), str(device))
+                self.__class__.__name__,
+                True,
+                total_bytes / (1024**2),
+                str(device),
+            )
         except Exception:
             pass
         return loaded_weights
-

vllm_omni/model_executor/models/qwen2_old.py

@@ -11,27 +11,38 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
-                                               QKVParallelLinear,
-                                               RowParallelLinear)
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.pooler import Pooler, PoolingType
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
-from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
 from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead, VocabParallelEmbedding)
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
 from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader, maybe_remap_kv_scale_name)
-from vllm.v1.pool.metadata import PoolingMetadata
-from vllm.model_executor.sampling_metadata import SamplingMetadata
-from vllm.sequence import IntermediateTensors, PoolerOutput
-
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
 from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
-from vllm.model_executor.models.utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,
-                    is_pp_missing_parameter,
-                    make_empty_intermediate_tensors_factory, make_layers,
-                    maybe_prefix)
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import PoolerOutput, SamplerOutput
+from vllm.v1.pool.metadata import PoolingMetadata
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
 
 logger = init_logger(__name__)
 
@@ -62,8 +73,10 @@ class Qwen2MLP(nn.Module):
             prefix=f"{prefix}.down_proj",
         )
         if hidden_act != "silu":
-            raise ValueError(f"Unsupported activation: {hidden_act}. "
-                             "Only silu is supported for now.")
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. "
+                "Only silu is supported for now."
+            )
         self.act_fn = SiluAndMul()
 
     def forward(self, x):
@@ -75,18 +88,20 @@ class Qwen2MLP(nn.Module):
 
 class Qwen2Attention(nn.Module):
 
-    def __init__(self,
-                 hidden_size: int,
-                 num_heads: int,
-                 num_kv_heads: int,
-                 max_position: int = 4096 * 32,
-                 rope_theta: float = 10000,
-                 head_dim: Optional[int] = None,
-                 cache_config: Optional[CacheConfig] = None,
-                 quant_config: Optional[QuantizationConfig] = None,
-                 rope_scaling: Optional[Tuple] = None,
-                 prefix: str = "",
-                 attn_type: str = AttentionType.DECODER) -> None:
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        max_position: int = 4096 * 32,
+        rope_theta: float = 10000,
+        head_dim: Optional[int] = None,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        rope_scaling: Optional[Tuple] = None,
+        prefix: str = "",
+        attn_type: str = AttentionType.DECODER,
+    ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
         tp_size = get_tensor_model_parallel_world_size()
@@ -133,14 +148,16 @@ class Qwen2Attention(nn.Module):
             base=self.rope_theta,
             rope_scaling=rope_scaling,
         )
-        self.attn = Attention(self.num_heads,
-                              self.head_dim,
-                              self.scaling,
-                              num_kv_heads=self.num_kv_heads,
-                              cache_config=cache_config,
-                              quant_config=quant_config,
-                              prefix=f"{prefix}.attn",
-                              attn_type=attn_type)
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+            attn_type=attn_type,
+        )
 
     def forward(
         self,
@@ -199,10 +216,10 @@ class Qwen2DecoderLayer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.mlp",
         )
-        self.input_layernorm = RMSNorm(config.hidden_size,
-                                       eps=config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(config.hidden_size,
-                                                eps=config.rms_norm_eps)
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
 
     def forward(
         self,
@@ -215,16 +232,14 @@ class Qwen2DecoderLayer(nn.Module):
             residual = hidden_states
             hidden_states = self.input_layernorm(hidden_states)
         else:
-            hidden_states, residual = self.input_layernorm(
-                hidden_states, residual)
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,
         )
 
         # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(
-            hidden_states, residual)
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
         hidden_states = self.mlp(hidden_states)
         return hidden_states, residual
 
@@ -237,14 +252,17 @@ class Qwen2DecoderLayer(nn.Module):
         "positions": -1,
         "intermediate_tensors": 0,
         "inputs_embeds": 0,
-    })
+    }
+)
 class Qwen2Model(nn.Module):
 
-    def __init__(self,
-                 *,
-                 vllm_config: VllmConfig,
-                 prefix: str = "",
-                 decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+        decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer,
+    ):
         super().__init__()
 
         config = vllm_config.model_config.hf_config
@@ -252,23 +270,27 @@ class Qwen2Model(nn.Module):
         quant_config = vllm_config.quant_config
 
         # TODO (@robertgshaw2): see if this can be moved out
-        if (cache_config.sliding_window is not None
-                and hasattr(config, "max_window_layers")):
-            raise ValueError("Sliding window for some but all layers is not "
-                             "supported. This model uses sliding window "
-                             "but `max_window_layers` = {} is less than "
-                             "`num_hidden_layers` = {}. Please open an issue "
-                             "to discuss this feature.".format(
-                                 config.max_window_layers,
-                                 config.num_hidden_layers,
-                             ))
+        if cache_config.sliding_window is not None and hasattr(
+            config, "max_window_layers"
+        ):
+            raise ValueError(
+                "Sliding window for some but all layers is not "
+                "supported. This model uses sliding window "
+                "but `max_window_layers` = {} is less than "
+                "`num_hidden_layers` = {}. Please open an issue "
+                "to discuss this feature.".format(
+                    config.max_window_layers,
+                    config.num_hidden_layers,
+                )
+            )
 
         self.config = config
         self.quant_config = quant_config
         self.vocab_size = config.vocab_size
 
-        if get_pp_group().is_first_rank or (config.tie_word_embeddings
-                                            and get_pp_group().is_last_rank):
+        if get_pp_group().is_first_rank or (
+            config.tie_word_embeddings and get_pp_group().is_last_rank
+        ):
             self.embed_tokens = VocabParallelEmbedding(
                 config.vocab_size,
                 getattr(config, "embedding_size", config.hidden_size),
@@ -282,16 +304,18 @@ class Qwen2Model(nn.Module):
         decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
-            lambda prefix: decoder_layer_type(config=config,
-                                              cache_config=cache_config,
-                                              quant_config=quant_config,
-                                              prefix=prefix),
+            lambda prefix: decoder_layer_type(
+                config=config,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=prefix,
+            ),
             prefix=f"{prefix}.layers",
         )
 
-        self.make_empty_intermediate_tensors = (
-            make_empty_intermediate_tensors_factory(
-                ["hidden_states", "residual"], config.hidden_size))
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
         if get_pp_group().is_last_rank:
             self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         else:
@@ -317,22 +341,20 @@ class Qwen2Model(nn.Module):
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
-        for layer in self.layers[self.start_layer:self.end_layer]:
+        for layer in self.layers[self.start_layer : self.end_layer]:
             hidden_states, residual = layer(
                 positions,
                 hidden_states,
                 residual,
             )
         if not get_pp_group().is_last_rank:
-            return IntermediateTensors({
-                "hidden_states": hidden_states,
-                "residual": residual
-            })
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
@@ -346,18 +368,19 @@ class Qwen2Model(nn.Module):
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
-            if (self.quant_config is not None and
-                (scale_name := self.quant_config.get_cache_scale(name))):
+            if self.quant_config is not None and (
+                scale_name := self.quant_config.get_cache_scale(name)
+            ):
                 # Loading kv cache quantization scales
                 param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader",
-                                        default_weight_loader)
-                loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
-                                 loaded_weight[0])
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                loaded_weight = (
+                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
+                )
                 weight_loader(param, loaded_weight)
                 loaded_params.add(scale_name)
                 continue
-            for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
                 name = name.replace(weight_name, param_name)
@@ -381,8 +404,7 @@ class Qwen2Model(nn.Module):
                 if is_pp_missing_parameter(name, self):
                     continue
                 param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader",
-                                        default_weight_loader)
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                 weight_loader(param, loaded_weight)
             loaded_params.add(name)
         return loaded_params
@@ -411,26 +433,29 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         self.lora_config = lora_config
 
         self.quant_config = quant_config
-        self.model = Qwen2Model(vllm_config=vllm_config,
-                                prefix=maybe_prefix(prefix, "model"))
+        self.model = Qwen2Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
 
         if get_pp_group().is_last_rank:
             if config.tie_word_embeddings:
                 self.lm_head = self.model.embed_tokens
             else:
-                self.lm_head = ParallelLMHead(config.vocab_size,
-                                              config.hidden_size,
-                                              quant_config=quant_config,
-                                              prefix=maybe_prefix(
-                                                  prefix, "lm_head"))
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
         else:
             self.lm_head = PPMissingLayer()
 
         self.logits_processor = LogitsProcessor(config.vocab_size)
-        self.sampler = get_sampler()
+        self.sampler = Sampler()
 
         self.make_empty_intermediate_tensors = (
-            self.model.make_empty_intermediate_tensors)
+            self.model.make_empty_intermediate_tensors
+        )
 
     def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.get_input_embeddings(input_ids)
@@ -442,17 +467,16 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         intermediate_tensors: Optional[IntermediateTensors] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, IntermediateTensors]:
-        hidden_states = self.model(input_ids, positions, intermediate_tensors,
-                                   inputs_embeds)
+        hidden_states = self.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )
         return hidden_states
 
     def compute_logits(
         self,
         hidden_states: torch.Tensor,
-        sampling_metadata: SamplingMetadata,
     ) -> Optional[torch.Tensor]:
-        logits = self.logits_processor(self.lm_head, hidden_states,
-                                       sampling_metadata)
+        logits = self.logits_processor(self.lm_head, hidden_states)
         return logits
 
     def sample(
@@ -463,12 +487,10 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         next_tokens = self.sampler(logits, sampling_metadata)
         return next_tokens
 
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         loader = AutoWeightsLoader(
             self,
-            skip_prefixes=(["lm_head."]
-                           if self.config.tie_word_embeddings else None),
+            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
         )
         return loader.load_weights(weights)
 
@@ -499,8 +521,9 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
         self.lora_config = lora_config
 
         self.quant_config = quant_config
-        self.model = Qwen2Model(vllm_config=vllm_config,
-                                prefix=maybe_prefix(prefix, "model"))
+        self.model = Qwen2Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
 
         # TODO: Replace this model class with as_embedding_model(
         # Qwen2ForCausalLM) after changing the default pooling method
@@ -508,14 +531,13 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
             logger.warning(
                 "This embedding model will default to last-token pooling in "
                 "an upcoming version. To avoid breaking changes, you should "
-                "pass `--override-pooler-config '{\"pooling_type\": \"MEAN\"}'`"
-                " explicitly.")
+                'pass `--override-pooler-config \'{"pooling_type": "MEAN"}\'`'
+                " explicitly."
+            )
 
         self._pooler = Pooler.from_config_with_defaults(
-            pooler_config,
-            pooling_type=PoolingType.MEAN,
-            normalize=True,
-            softmax=False)
+            pooler_config, pooling_type=PoolingType.MEAN, normalize=True, softmax=False
+        )
 
     def forward(
         self,
@@ -534,6 +556,7 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         weights = self.hf_to_vllm_mapper.apply(weights)
-        weights = ((name, data) for name, data in weights
-                   if not name.startswith("lm_head."))
-        self.model.load_weights(weights)
+        weights = (
+            (name, data) for name, data in weights if not name.startswith("lm_head.")
+        )
+        self.model.load_weights(weights)

vllm_omni/model_executor/models/registry.py

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

vllm_omni/model_executor/models/utils.py

@@ -5,4 +5,4 @@ def add_prefix_to_loaded_weights(weights: set[str], prefix: str) -> set[str]:
     """
     Add a prefix to the names of the loaded weights.
     """
-    return {maybe_prefix(prefix, name) for name in weights}
+    return {maybe_prefix(prefix, name) for name in weights}

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -1,40 +1,68 @@
 # stage config for running qwen2.5-omni with architecture of OmniLLM.
 stage_args:
   - stage_id: 0
+    runtime:
+      process: true            # Run this stage in a separate process
+      devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+      max_batch_size: 1
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_cls: vllm_omni.worker.gpu.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-      gpu_memory_utilization: 0.32
+      gpu_memory_utilization: 0.8
       enforce_eager: true  # need to discuss
       trust_remote_code: true
       engine_output_type: latent  # change the param name,such as pooling_output
+      enable_prefix_caching: false
     final_output: true
     final_output_type: text
   - stage_id: 1
+    runtime:
+      process: true
+      devices: "1"
+      max_batch_size: 3
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_cls: vllm_omni.worker.gpu.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-      gpu_memory_utilization: 0.32
+      gpu_memory_utilization: 0.8
       enforce_eager: true
       trust_remote_code: true
+      enable_prefix_caching: false
       engine_output_type: latent
     engine_input_source: [0]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
   - stage_id: 2
+    runtime:
+      process: true
+      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
+      worker_cls: vllm_omni.worker.gpu.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
-      gpu_memory_utilization: 0.3
+      gpu_memory_utilization: 0.15
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false
       engine_output_type: audio
     engine_input_source: [1]
     final_output: true
-    final_output_type: audio
+    final_output_type: audio
+
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1             # Simplified: trigger downstream only after full upstream completion
+    max_inflight: 1             # Simplified: process serially within each stage
+  edges:
+    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
+      to: 1
+      window_size: -1
+    - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
+      to: 2
+      window_size: -1

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -1,40 +1,58 @@
 from typing import Union
+
 import torch
 
 from vllm.inputs import TextPrompt
-
 from vllm_omni.inputs.data import OmniTokensPrompt
-
-def thinker2talker(stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None):
+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
+):
     if not engine_input_source:
-      raise ValueError("engine_input_source cannot be empty")
+        raise ValueError("engine_input_source cannot be empty")
     source_stage_id = engine_input_source[0]
     if source_stage_id >= len(stage_list):
-      raise IndexError(f"Invalid stage_id: {source_stage_id}")
+        raise IndexError(f"Invalid stage_id: {source_stage_id}")
     if stage_list[source_stage_id].engine_outputs is None:
         raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
     thinker_outputs = stage_list[source_stage_id].engine_outputs
     talker_inputs = []
-    multi_modal_data = {thinker_output.request_id: 
-            prompt.get('multi_modal_data', None) for thinker_output, prompt in zip(thinker_outputs, prompt)}
-    
+    multi_modal_data = {
+        thinker_output.request_id: prompt.get("multi_modal_data", None)
+        for thinker_output, prompt in zip(thinker_outputs, prompt)
+    }
+
     for i, thinker_output in enumerate(thinker_outputs):
         output = thinker_output.outputs[0]
         prompt_token_ids = thinker_output.prompt_token_ids
         thinker_output_ids = output.token_ids
         prompt_token_ids_len = len(prompt_token_ids)
-        thinker_hidden_states = output.multimodal_output["latent"].clone().detach().cuda()
+        thinker_hidden_states = (
+            output.multimodal_output["latent"].clone().detach().cuda()
+        )
         talker_inputs.append(
             OmniTokensPrompt(
-                prompt_token_ids=[0] * (len(prompt_token_ids) + 2), # add 2 for codec pad and start token
+                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),
+                    "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),
                 },
-                multi_modal_data=multi_modal_data[thinker_output.request_id] if multi_modal_data is not None else None,
+                multi_modal_data=(
+                    multi_modal_data[thinker_output.request_id]
+                    if multi_modal_data is not None
+                    else None
+                ),
                 mm_processor_kwargs=None,
             )
         )
-    return talker_inputs
+    return talker_inputs

vllm_omni/outputs.py

@@ -1,9 +1,18 @@
-from vllm.outputs import RequestOutput
 from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from vllm.outputs import RequestOutput
+from vllm.v1.outputs import ModelRunnerOutput
+
+
+class OmniModelRunnerOutput(ModelRunnerOutput):
+    multimodal_outputs: Optional[dict[str, torch.Tensor]] = None
 
 
 @dataclass
 class OmniRequestOutput(RequestOutput):
     stage_id: int
     final_output_type: str
-    request_output: RequestOutput
+    request_output: RequestOutput

vllm_omni/patch.py

@@ -0,0 +1,28 @@
+import sys
+
+from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
+from vllm.model_executor.layers.rotary_embedding import (
+    MRotaryEmbedding as _OriginalMRotaryEmbedding,
+)
+from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest
+from vllm.v1.request import Request as _OriginalRequest
+from vllm_omni.engine import OmniEngineCoreRequest
+from vllm_omni.inputs.data import OmniTokensPrompt
+from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
+from vllm_omni.request import OmniRequest
+
+for module_name, module in sys.modules.items():
+    if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt:
+        module.TokensPrompt = OmniTokensPrompt
+    if (
+        hasattr(module, "MRotaryEmbedding")
+        and module.MRotaryEmbedding == _OriginalMRotaryEmbedding
+    ):
+        module.MRotaryEmbedding = MRotaryEmbedding
+    if hasattr(module, "Request") and module.Request == _OriginalRequest:
+        module.Request = OmniRequest
+    if (
+        hasattr(module, "EngineCoreRequest")
+        and module.EngineCoreRequest == _OriginalEngineCoreRequest
+    ):
+        module.EngineCoreRequest = OmniEngineCoreRequest

vllm_omni/request.py

@@ -1,299 +1,60 @@
-"""
-OmniRequest: Extended request class for vLLM-omni multimodal processing.
+from typing import TYPE_CHECKING, Callable, Optional
 
-This class extends vLLM's Request to support multimodal and non-autoregressive
-processing with additional fields and methods specific to vLLM-omni.
-"""
+from vllm.v1.request import Request
+from vllm.v1.structured_output.request import StructuredOutputRequest
 
-import enum
-import time
-from typing import Optional, Dict, Any, List, Union
-from dataclasses import dataclass, field
-from vllm.v1.request import Request as vLLMRequest
+if TYPE_CHECKING:
+    from vllm.v1.core.kv_cache_utils import BlockHash
+
+from vllm_omni.engine import (
+    AdditionalInformationPayload,
+    OmniEngineCoreRequest,
+    PromptEmbedsPayload,
+)
 
 
-class RequestType(enum.Enum):
-    """Types of requests supported by vLLM-omni."""
-    TEXT = "text"
-    IMAGE = "image"
-    AUDIO = "audio"
-    VIDEO = "video"
-    MULTIMODAL = "multimodal"
-    DIFFUSION = "diffusion"
-
-
-class ProcessingStage(enum.Enum):
-    """Processing stages for multi-stage models."""
-    PREPROCESSING = "preprocessing"
-    AR_GENERATION = "ar_generation"
-    DIFFUSION_GENERATION = "diffusion_generation"
-    POSTPROCESSING = "postprocessing"
-    COMPLETED = "completed"
-
-
-@dataclass
-class MultimodalData:
-    """Container for multimodal input data."""
-    data_type: str  # "image", "audio", "video", etc.
-    data: Any  # The actual data (numpy array, bytes, etc.)
-    metadata: Dict[str, Any] = field(default_factory=dict)
-    processed: bool = False
-
-
-@dataclass
-class DiffusionParams:
-    """Parameters specific to diffusion model generation."""
-    num_inference_steps: int = 50
-    guidance_scale: float = 7.5
-    seed: Optional[int] = None
-    scheduler: str = "ddpm"
-    strength: float = 1.0
-    eta: float = 0.0
-
-
-class OmniRequest(vLLMRequest):
-    """
-    Extended request class for vLLM-omni multimodal and non-autoregressive processing.
-    
-    This class extends vLLM's Request with additional fields and methods to support:
-    - Multimodal input processing
-    - Non-autoregressive generation (diffusion models)
-    - Multi-stage processing pipelines
-    - Enhanced caching for different model types
-    """
-    
+class OmniRequest(Request):
     def __init__(
         self,
-        request_id: str,
-        prompt: Optional[str] = None,
-        prompt_token_ids: Optional[List[int]] = None,
-        sampling_params: Optional[Any] = None,
-        arrival_time: Optional[float] = None,
-        lora_request: Optional[Any] = None,
-        multi_modal_data: Optional[Dict[str, Any]] = None,
-        multi_modal_placeholders: Optional[Dict[str, str]] = None,
-        priority: int = 0,
-        # vLLM-omni specific parameters
-        request_type: RequestType = RequestType.TEXT,
-        processing_stage: ProcessingStage = ProcessingStage.PREPROCESSING,
-        multimodal_inputs: Optional[List[MultimodalData]] = None,
-        diffusion_params: Optional[DiffusionParams] = None,
-        output_format: str = "text",
-        cache_key: Optional[str] = None,
-        stage_configs: Optional[Dict[str, Any]] = None,
-        **kwargs
+        prompt_embeds: Optional[PromptEmbedsPayload] = None,
+        additional_information: Optional[AdditionalInformationPayload] = None,
+        *args,
+        **kwargs,
     ):
-        # Initialize parent class
-        super().__init__(
-            request_id=request_id,
-            prompt=prompt,
-            prompt_token_ids=prompt_token_ids,
-            sampling_params=sampling_params,
-            arrival_time=arrival_time or time.time(),
-            lora_request=lora_request,
-            multi_modal_data=multi_modal_data,
-            multi_modal_placeholders=multi_modal_placeholders,
-            priority=priority,
-            **kwargs
+        super().__init__(*args, **kwargs)
+        # Serialized prompt embeddings payload (optional)
+        self.prompt_embeds: Optional[PromptEmbedsPayload] = prompt_embeds
+        # Serialized additional information payload (optional)
+        self.additional_information: Optional[AdditionalInformationPayload] = (
+            additional_information
         )
-        
-        # vLLM-omni specific attributes
-        self.request_type = request_type
-        self.processing_stage = processing_stage
-        self.multimodal_inputs = multimodal_inputs or []
-        self.diffusion_params = diffusion_params or DiffusionParams()
-        self.output_format = output_format
-        self.cache_key = cache_key
-        self.stage_configs = stage_configs or {}
-        
-        # Processing state
-        self.current_stage = 0
-        self.stage_results = {}
-        self.hidden_states = None
-        self.intermediate_outputs = []
-        
-        # Timing and metrics
-        self.stage_timings = {}
-        self.total_processing_time = 0.0
-        
-        # Error handling
-        self.errors = []
-        self.retry_count = 0
-        self.max_retries = 3
-    
-    def add_multimodal_input(self, data: MultimodalData) -> None:
-        """Add a multimodal input to the request."""
-        self.multimodal_inputs.append(data)
-    
-    def get_multimodal_inputs_by_type(self, data_type: str) -> List[MultimodalData]:
-        """Get all multimodal inputs of a specific type."""
-        return [inp for inp in self.multimodal_inputs if inp.data_type == data_type]
-    
-    def update_processing_stage(self, stage: ProcessingStage) -> None:
-        """Update the current processing stage."""
-        self.processing_stage = stage
-        self.stage_timings[stage.value] = time.time()
-    
-    def add_stage_result(self, stage: str, result: Any) -> None:
-        """Add a result from a processing stage."""
-        self.stage_results[stage] = result
-        self.intermediate_outputs.append({
-            'stage': stage,
-            'result': result,
-            'timestamp': time.time()
-        })
-    
-    def get_stage_result(self, stage: str) -> Any:
-        """Get a result from a specific processing stage."""
-        return self.stage_results.get(stage)
-    
-    def set_hidden_states(self, hidden_states: Any) -> None:
-        """Set hidden states for the request."""
-        self.hidden_states = hidden_states
-    
-    def get_hidden_states(self) -> Any:
-        """Get hidden states for the request."""
-        return self.hidden_states
-    
-    def add_error(self, error: str) -> None:
-        """Add an error to the request."""
-        self.errors.append({
-            'error': error,
-            'timestamp': time.time(),
-            'stage': self.processing_stage.value
-        })
-    
-    def has_errors(self) -> bool:
-        """Check if the request has any errors."""
-        return len(self.errors) > 0
-    
-    def can_retry(self) -> bool:
-        """Check if the request can be retried."""
-        return self.retry_count < self.max_retries
-    
-    def increment_retry(self) -> None:
-        """Increment the retry count."""
-        self.retry_count += 1
-    
-    def generate_cache_key(self) -> str:
-        """Generate a cache key for the request."""
-        if self.cache_key:
-            return self.cache_key
-        
-        # Generate cache key based on request content
-        key_parts = [
-            self.request_id,
-            str(self.request_type.value),
-            str(self.prompt_token_ids) if self.prompt_token_ids else str(self.prompt),
-            str(self.diffusion_params.num_inference_steps) if self.diffusion_params else "0",
-            str(len(self.multimodal_inputs))
-        ]
-        return "_".join(key_parts)
-    
-    def get_processing_time(self) -> float:
-        """Get the total processing time for the request."""
-        if self.stage_timings:
-            return time.time() - min(self.stage_timings.values())
-        return 0.0
-    
-    def is_completed(self) -> bool:
-        """Check if the request is completed."""
-        return self.processing_stage == ProcessingStage.COMPLETED
-    
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert the request to a dictionary for serialization."""
-        return {
-            'request_id': self.request_id,
-            'request_type': self.request_type.value,
-            'processing_stage': self.processing_stage.value,
-            'prompt': self.prompt,
-            'prompt_token_ids': self.prompt_token_ids,
-            'output_format': self.output_format,
-            'multimodal_inputs_count': len(self.multimodal_inputs),
-            'diffusion_params': self.diffusion_params.__dict__ if self.diffusion_params else None,
-            'stage_results': list(self.stage_results.keys()),
-            'has_errors': self.has_errors(),
-            'retry_count': self.retry_count,
-            'processing_time': self.get_processing_time()
-        }
-    
-    def __repr__(self) -> str:
-        """String representation of the request."""
-        return (f"OmniRequest(id={self.request_id}, "
-                f"type={self.request_type.value}, "
-                f"stage={self.processing_stage.value}, "
-                f"multimodal_inputs={len(self.multimodal_inputs)})")
 
+    @classmethod
+    def from_engine_core_request(
+        cls,
+        request: OmniEngineCoreRequest,
+        block_hasher: Optional[Callable[["Request"], list["BlockHash"]]],
+    ) -> "Request":
 
-# Factory functions for creating different types of requests
-def create_text_request(
-    request_id: str,
-    prompt: str,
-    sampling_params: Optional[Any] = None,
-    **kwargs
-) -> OmniRequest:
-    """Create a text-only request."""
-    return OmniRequest(
-        request_id=request_id,
-        prompt=prompt,
-        request_type=RequestType.TEXT,
-        sampling_params=sampling_params,
-        **kwargs
-    )
-
-
-def create_image_request(
-    request_id: str,
-    prompt: str,
-    image_data: Any,
-    diffusion_params: Optional[DiffusionParams] = None,
-    **kwargs
-) -> OmniRequest:
-    """Create an image generation request."""
-    multimodal_input = MultimodalData(
-        data_type="image",
-        data=image_data,
-        metadata={"is_input": True}
-    )
-    
-    return OmniRequest(
-        request_id=request_id,
-        prompt=prompt,
-        request_type=RequestType.IMAGE,
-        multimodal_inputs=[multimodal_input],
-        diffusion_params=diffusion_params,
-        output_format="image",
-        **kwargs
-    )
-
-
-def create_multimodal_request(
-    request_id: str,
-    prompt: str,
-    multimodal_inputs: List[MultimodalData],
-    **kwargs
-) -> OmniRequest:
-    """Create a multimodal request."""
-    return OmniRequest(
-        request_id=request_id,
-        prompt=prompt,
-        request_type=RequestType.MULTIMODAL,
-        multimodal_inputs=multimodal_inputs,
-        **kwargs
-    )
-
-
-def create_diffusion_request(
-    request_id: str,
-    prompt: str,
-    diffusion_params: DiffusionParams,
-    **kwargs
-) -> OmniRequest:
-    """Create a diffusion model request."""
-    return OmniRequest(
-        request_id=request_id,
-        prompt=prompt,
-        request_type=RequestType.DIFFUSION,
-        diffusion_params=diffusion_params,
-        **kwargs
-    )
+        return cls(
+            request_id=request.request_id,
+            client_index=request.client_index,
+            prompt_token_ids=request.prompt_token_ids,
+            mm_features=request.mm_features,
+            sampling_params=request.sampling_params,
+            pooling_params=request.pooling_params,
+            eos_token_id=request.eos_token_id,
+            arrival_time=request.arrival_time,
+            lora_request=request.lora_request,
+            structured_output_request=(
+                StructuredOutputRequest(sampling_params=request.sampling_params)
+                if request.sampling_params
+                else None
+            ),
+            cache_salt=request.cache_salt,
+            priority=request.priority,
+            trace_headers=request.trace_headers,
+            block_hasher=block_hasher,
+            prompt_embeds=request.prompt_embeds,
+            additional_information=request.additional_information,
+        )

vllm_omni/worker/cpu/cpu_ar_model_runner.py

@@ -0,0 +1,649 @@
+"""AR CPU Model Runner for vLLM-omni.
+
+Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
+and also outputs sampled tokens.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Union
+
+import numpy as np
+import torch
+
+from vllm.forward_context import BatchDescriptor
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import AsyncModelRunnerOutput
+from vllm.v1.structured_output.utils import apply_grammar_bitmask
+from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker.gpu_model_runner import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    AsyncGPUModelRunnerOutput,
+    IntermediateTensors,
+    get_pp_group,
+    get_tp_group,
+    has_kv_transfer_group,
+    set_forward_context,
+)
+from vllm.v1.worker.ubatch_utils import UBatchSlices
+from vllm.v1.worker.utils import is_residual_scattered_for_sp
+from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.cpu.cpu_model_runner import OmniCPUModelRunner
+
+logger = init_logger(__name__)
+
+
+class CPUARModelRunner(OmniCPUModelRunner):
+    """Autoregressive GPU model runner that returns hidden states per request.
+
+    This runner follows the same preparation and forward path as GPUModelRunner
+    (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs),
+    and additionally performs lightweight sampling so that sampled tokens are
+    available in outputs. Hidden representations are taken at the same indices
+    that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`).
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
+
+    def _preprocess(
+        self,
+        scheduler_output: "SchedulerOutput",
+        num_scheduled_tokens_np: np.ndarray,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        ubatch_slices: Optional[UBatchSlices] = None,
+        num_tokens_after_padding: Optional[torch.Tensor] = None,
+    ) -> tuple[
+        int,
+        int,
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        torch.Tensor,
+        Optional[IntermediateTensors],
+        dict[str, Any],
+        Optional[dict[str, dict]],
+    ]:
+
+        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        if ubatch_slices:
+            assert num_tokens_after_padding is not None
+            num_input_tokens = int(num_tokens_after_padding[0].item() * 2)
+            self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
+        elif ubatch_slices is None:
+            num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)
+            num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
+            num_input_tokens += num_pad
+
+        # _prepare_inputs may reorder the batch, so we must gather multi
+        # modal outputs after that to ensure the correct order
+        per_req_additional_information: Optional[dict[str, dict]] = None
+        if (
+            self.supports_mm_inputs
+            and get_pp_group().is_first_rank
+            and not self.model_config.is_encoder_decoder
+        ):
+            # Run the multimodal encoder if any.
+            self._execute_mm_encoder(scheduler_output)
+            mm_embeds = self._gather_mm_embeddings(scheduler_output)
+
+            # NOTE(woosuk): To unify token ids and soft tokens (vision
+            # embeddings), we always use embeddings (rather than token ids)
+            # as input to the multimodal model, even when the input is text.
+            inputs_embeds_scheduled = self.model.get_input_embeddings(
+                input_ids=self.input_ids.gpu[:num_scheduled_tokens],
+                multimodal_embeddings=mm_embeds or None,
+            )
+
+            # TODO(woosuk): Avoid the copy. Optimize.
+            self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
+
+            # Reset per-step additional information collector (deprecated concat path)
+            if hasattr(self, "_forward_additional_information"):
+                self._forward_additional_information = None
+            # New: per-request additional information for this step
+            per_req_additional_information = {}
+
+            # Overlay custom prompt_embeds per request for the prompt portion;
+            # collect additional_information (tensor/list) for prefill portion only
+            for req_index, req_id in enumerate(self.input_batch.req_ids):
+                req_state = self.requests[req_id]
+                pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
+                addi_cpu = getattr(req_state, "additional_information_cpu", None)
+                num_computed_tokens = int(
+                    self.input_batch.num_computed_tokens_cpu[req_index]
+                )
+                prompt_len = len(req_state.prompt_token_ids)
+                prompt_remaining = max(0, prompt_len - num_computed_tokens)
+                sched_tokens = int(num_scheduled_tokens_np[req_index])
+                overlay_len = min(sched_tokens, prompt_remaining)
+                if overlay_len <= 0:
+                    continue
+                if pe_cpu is not None:
+                    src = pe_cpu[
+                        num_computed_tokens : num_computed_tokens + overlay_len
+                    ].to(dtype=self.dtype, device=self.device, non_blocking=True)
+                    start_offset = int(self.query_start_loc.cpu[req_index])
+                    self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(
+                        src
+                    )
+                # Build per-request additional information (no cross-request concat)
+                if addi_cpu is not None and isinstance(addi_cpu, dict):
+                    req_info: dict[str, object] = {}
+                    for k, v in addi_cpu.items():
+                        if isinstance(v, torch.Tensor):
+                            # 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()
+                                except Exception:
+                                    seg = v.detach().to("cpu").contiguous()
+                                req_info[k] = seg
+                            else:
+                                req_info[k] = v.detach().to("cpu").contiguous()
+                        elif isinstance(v, list):
+                            req_info[k] = v
+                        else:
+                            req_info[k] = v
+                    per_req_additional_information[req_id] = req_info
+
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = {
+                **self._init_model_kwargs(num_scheduled_tokens),
+                **self._extract_mm_kwargs(scheduler_output),
+            }
+        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
+            # Get the input embeddings for the tokens that are not input embeds,
+            # then put them into the appropriate positions.
+            # TODO(qthequartermasterman): Since even when prompt embeds are
+            # enabled, (a) not all requests will use prompt embeds, and (b)
+            # after the initial prompt is processed, the rest of the generated
+            # tokens will be token ids, it is not desirable to have the
+            # embedding layer outside of the CUDA graph all the time. The v0
+            # engine avoids this by "double compiling" the CUDA graph, once
+            # with input_ids and again with inputs_embeds, for all num_tokens.
+            # If a batch only has token ids, then including the embedding layer
+            # in the CUDA graph will be more performant (like in the else case
+            # below).
+            token_ids_idx = (
+                self.is_token_ids.gpu[:num_scheduled_tokens]
+                .nonzero(as_tuple=False)
+                .squeeze(1)
+            )
+            # Some tokens ids may need to become embeds
+            if token_ids_idx.numel() > 0:
+                token_ids = self.input_ids.gpu[token_ids_idx]
+                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
+                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
+
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+            input_ids = None
+        else:
+            # For text-only models, we use token ids as input.
+            # While it is possible to use embeddings as input just like the
+            # multimodal models, it is not desirable for performance since
+            # then the embedding layer is not included in the CUDA graph.
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = None
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+        if self.uses_mrope:
+            positions = self.mrope_positions.gpu[:, :num_input_tokens]
+        else:
+            positions = self.positions.gpu[:num_input_tokens]
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True
+            )
+
+        if (
+            self.model_config.is_encoder_decoder
+            and scheduler_output.scheduled_encoder_inputs
+        ):
+            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
+            model_kwargs.update(encoder_inputs)
+
+        return (
+            num_scheduled_tokens,
+            num_input_tokens,
+            num_tokens_after_padding,
+            input_ids,
+            inputs_embeds,
+            positions,
+            intermediate_tensors,
+            model_kwargs,
+            per_req_additional_information,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[OmniModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:
+        with record_function_or_nullcontext("Preprocess"):
+            with self.synchronize_input_prep():
+                # Update persistent batch states.
+                self._update_states(scheduler_output)
+
+                # Decode per-request prompt_embeds / additional_hidden_states payloads
+                # (if present) into CPU tensors
+                try:
+                    new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
+                    if new_reqs:
+                        for nr in new_reqs:
+                            req_id = getattr(nr, "req_id", None) or getattr(
+                                nr, "request_id", None
+                            )
+                            if req_id is None:
+                                continue
+                            # prompt_embeds
+                            payload_pe = getattr(nr, "prompt_embeds", None)
+                            if payload_pe is not None:
+                                if isinstance(payload_pe, torch.Tensor):
+                                    pe_cpu = payload_pe.detach().to("cpu").contiguous()
+                                elif isinstance(payload_pe, PromptEmbedsPayload):
+                                    dt = np.dtype(
+                                        getattr(payload_pe, "dtype", "float32")
+                                    )
+                                    arr = np.frombuffer(payload_pe.data, dtype=dt)
+                                    arr = arr.reshape(payload_pe.shape)
+                                    pe_cpu = torch.from_numpy(arr)
+                                else:
+                                    pe_cpu = None
+                                if pe_cpu is not None and req_id in self.requests:
+                                    setattr(
+                                        self.requests[req_id],
+                                        "prompt_embeds_cpu",
+                                        pe_cpu,
+                                    )
+                            # additional_information
+                            payload_info = getattr(nr, "additional_information", None)
+                            if payload_info is not None:
+                                info_dict = {}
+                                if isinstance(payload_info, dict):
+                                    # Already decoded
+                                    info_dict = payload_info
+                                elif isinstance(
+                                    payload_info, AdditionalInformationPayload
+                                ):
+                                    for k, entry in payload_info.entries.items():
+                                        if entry.tensor_data is not None:
+                                            dt = np.dtype(
+                                                getattr(
+                                                    entry, "tensor_dtype", "float32"
+                                                )
+                                            )
+                                            arr = np.frombuffer(
+                                                entry.tensor_data, dtype=dt
+                                            )
+                                            arr = arr.reshape(entry.tensor_shape)
+                                            info_dict[k] = torch.from_numpy(arr)
+                                        else:
+                                            info_dict[k] = entry.list_data
+                                if info_dict and req_id in self.requests:
+                                    setattr(
+                                        self.requests[req_id],
+                                        "additional_information_cpu",
+                                        info_dict,
+                                    )
+                except Exception as e:
+                    logger.error(
+                        f"Error decoding prompt_embeds / additional_information: {e}"
+                    )
+                    pass
+
+                if not scheduler_output.total_num_scheduled_tokens:
+                    if not has_kv_transfer_group():
+                        # Return empty ModelRunnerOutput if no work to do.
+                        return EMPTY_MODEL_RUNNER_OUTPUT
+                    return self.kv_connector_no_forward(
+                        scheduler_output, self.vllm_config
+                    )
+                if self.cache_config.kv_sharing_fast_prefill:
+                    assert not self.input_batch.num_prompt_logprobs, (
+                        "--kv-sharing-fast-prefill produces incorrect "
+                        "logprobs for prompt tokens, tokens, please disable "
+                        "it when the requests need prompt logprobs"
+                    )
+
+                # Prepare the decoder inputs.
+                (
+                    attn_metadata,
+                    logits_indices,
+                    spec_decode_metadata,
+                    num_scheduled_tokens_np,
+                    spec_decode_common_attn_metadata,
+                    max_query_len,
+                    ubatch_slices,
+                    num_tokens_after_padding,
+                ) = self._prepare_inputs(scheduler_output)
+
+            (
+                num_scheduled_tokens,
+                num_input_tokens,
+                num_tokens_across_dp,
+                input_ids,
+                inputs_embeds,
+                positions,
+                intermediate_tensors,
+                model_kwargs,
+                per_req_additional_information,
+            ) = self._preprocess(
+                scheduler_output,
+                num_scheduled_tokens_np,
+                intermediate_tensors,
+                ubatch_slices,
+                num_tokens_after_padding,
+            )
+
+            uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
+                num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
+            )
+            batch_descriptor = BatchDescriptor(
+                num_tokens=num_input_tokens, uniform_decode=uniform_decode
+            )
+            cudagraph_runtime_mode, batch_descriptor = (
+                self.cudagraph_dispatcher.dispatch(batch_descriptor)
+            )
+
+        # This is currently to get around the assert in the DPMetadata
+        # where it wants `num_tokens_across_dp` to align with `num_tokens`
+        if ubatch_slices is not None:
+            num_input_tokens = ubatch_slices[0].num_tokens
+
+        # Run the model.
+        # Use persistent buffers for CUDA graphs.
+        with (
+            set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
+                batch_descriptor=batch_descriptor,
+                ubatch_slices=ubatch_slices,
+            ),
+            record_function_or_nullcontext("Forward"),
+            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+        ):
+
+            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
+            # Always pass per-request runtime additional_information (persisted in request state)
+            try:
+                per_req_runtime_info = []
+                for req_id in self.input_batch.req_ids:
+                    req_state = self.requests.get(req_id)
+                    info = (
+                        getattr(req_state, "additional_information_cpu", None)
+                        if req_state is not None
+                        else None
+                    )
+                    per_req_runtime_info.append(info if isinstance(info, dict) else {})
+                model_kwargs_extra["runtime_additional_information"] = (
+                    per_req_runtime_info
+                )
+                model_kwargs_extra["request_ids"] = self.input_batch.req_ids
+                # Pass each request's token span within the flattened sequence for this step, enabling the model to map decode/prefill by request
+                req_token_spans = []
+                for req_index in range(len(self.input_batch.req_ids)):
+                    start_offset = int(self.query_start_loc.cpu[req_index])
+                    sched_tokens = int(num_scheduled_tokens_np[req_index])
+                    req_token_spans.append((start_offset, start_offset + sched_tokens))
+                model_kwargs_extra["request_token_spans"] = req_token_spans
+            except Exception:
+                pass
+            model_output = self.model(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+                **model_kwargs,
+                sampling_metadata=self.input_batch.sampling_metadata,
+                logits_index=logits_indices,
+                sampler=self.sampler,
+                **model_kwargs_extra,
+            )
+
+        with record_function_or_nullcontext("Postprocess"):
+            if self.use_aux_hidden_state_outputs:
+                # True when EAGLE 3 is used.
+                hidden_states, aux_hidden_states = model_output
+            else:
+                # Common case.
+                hidden_states = model_output
+                aux_hidden_states = None
+
+            hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+                hidden_states
+            )
+            # The model side may return per-request additional_information updates (model-agnostic channel).
+            # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
+            # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
+            try:
+                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(
+                        "additional_information_update"
+                    )
+                    if isinstance(updates_list, list):
+                        for idx, upd in enumerate(updates_list):
+                            if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
+                                continue
+                            req_id = self.input_batch.req_ids[idx]
+                            self._merge_additional_information_update(req_id, upd)
+                    # Option B: dict[str, dict] keyed by req_id
+                    updates_map = multimodal_outputs.get(
+                        "additional_information_update_by_req_id"
+                    )
+                    if isinstance(updates_map, dict):
+                        for req_id, upd in updates_map.items():
+                            if not isinstance(upd, dict):
+                                continue
+                            if req_id not in self.requests:
+                                continue
+                            self._merge_additional_information_update(req_id, upd)
+            except Exception as e:
+                logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}")
+            if not self.broadcast_pp_output:
+
+                # Common case.
+                if not get_pp_group().is_last_rank:
+                    # Return the intermediate tensors.
+                    assert isinstance(hidden_states, IntermediateTensors)
+                    hidden_states.kv_connector_output = kv_connector_output
+                    return hidden_states
+
+                if self.is_pooling_model:
+                    # Return the pooling output.
+                    output = self._pool(
+                        hidden_states, num_scheduled_tokens, num_scheduled_tokens_np
+                    )
+                    output.kv_connector_output = kv_connector_output
+                    return output
+
+                sample_hidden_states = hidden_states[logits_indices]
+                logits = self.model.compute_logits(sample_hidden_states)
+            else:
+                # Rare case.
+                assert not self.is_pooling_model
+
+                if not get_pp_group().is_last_rank:
+                    all_gather_tensors = {
+                        "residual": not is_residual_scattered_for_sp(
+                            self.vllm_config, num_input_tokens
+                        )
+                    }
+                    get_pp_group().send_tensor_dict(
+                        hidden_states.tensors,
+                        all_gather_group=get_tp_group(),
+                        all_gather_tensors=all_gather_tensors,
+                    )
+                    logits = None
+                else:
+                    sample_hidden_states = hidden_states[logits_indices]
+                    logits = self.model.compute_logits(sample_hidden_states)
+
+                model_output_broadcast_data = {}
+                if logits is not None:
+                    model_output_broadcast_data["logits"] = logits.contiguous()
+
+                model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
+                    model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
+                )
+                assert model_output_broadcast_data is not None
+                logits = model_output_broadcast_data["logits"]
+
+            # Apply structured output bitmasks if present
+            if scheduler_output.grammar_bitmask is not None:
+                apply_grammar_bitmask(
+                    scheduler_output, self.input_batch, logits, self.device
+                )
+
+        with record_function_or_nullcontext("Sample"):
+            sampler_output = self._sample(logits, spec_decode_metadata)
+
+        def propose_draft_token_ids(sampled_token_ids):
+            assert spec_decode_common_attn_metadata is not None
+            with record_function_or_nullcontext("Draft"):
+                self._draft_token_ids = self.propose_draft_token_ids(
+                    scheduler_output,
+                    sampled_token_ids,
+                    self.input_batch.sampling_metadata,
+                    hidden_states,
+                    sample_hidden_states,
+                    aux_hidden_states,
+                    spec_decode_metadata,
+                    spec_decode_common_attn_metadata,
+                )
+
+        use_padded_batch_for_eagle = (
+            self.speculative_config
+            and self.speculative_config.use_eagle()
+            and not self.speculative_config.disable_padded_drafter_batch
+        )
+        effective_drafter_max_model_len = self.max_model_len
+        if effective_drafter_max_model_len is None:
+            effective_drafter_max_model_len = self.model_config.max_model_len
+        if (
+            self.speculative_config
+            and self.speculative_config.draft_model_config is not None
+            and self.speculative_config.draft_model_config.max_model_len is not None
+        ):
+            effective_drafter_max_model_len = (
+                self.speculative_config.draft_model_config.max_model_len
+            )
+        input_fits_in_drafter = spec_decode_common_attn_metadata and (
+            spec_decode_common_attn_metadata.seq_lens.max()
+            + self.speculative_config.num_speculative_tokens
+            <= effective_drafter_max_model_len
+        )
+        if use_padded_batch_for_eagle and input_fits_in_drafter:
+            # EAGLE speculative decoding can use the GPU sampled tokens
+            # as inputs, and does not need to wait for bookkeeping to finish.
+            propose_draft_token_ids(sampler_output.sampled_token_ids)
+
+        with record_function_or_nullcontext("Bookkeep"):
+            (
+                num_nans_in_logits,
+                logprobs_lists,
+                valid_sampled_token_ids,
+                prompt_logprobs_dict,
+                req_ids_output_copy,
+                req_id_to_index_output_copy,
+                invalid_req_indices,
+            ) = self._bookkeeping_sync(
+                scheduler_output,
+                sampler_output,
+                logits,
+                hidden_states,
+                num_scheduled_tokens,
+            )
+
+        if (
+            self.speculative_config
+            and not use_padded_batch_for_eagle
+            and input_fits_in_drafter
+        ):
+            # ngram and other speculative decoding methods use the sampled
+            # tokens on the CPU, so they are run after bookkeeping.
+            propose_draft_token_ids(valid_sampled_token_ids)
+
+        with record_function_or_nullcontext("EPLB"):
+            self.eplb_step()
+
+        # Convert to per-request tensors on CPU
+        hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
+        pooler_output: list[Optional[torch.Tensor]] = []
+        prev_logits_index = 0
+        for logits_index in logits_indices:
+            pooler_output.append(
+                hidden_states_cpu[prev_logits_index : logits_index + 1]
+            )
+            prev_logits_index = logits_index + 1
+
+        output = OmniModelRunnerOutput(
+            req_ids=req_ids_output_copy,
+            req_id_to_index=req_id_to_index_output_copy,
+            sampled_token_ids=valid_sampled_token_ids,
+            logprobs=logprobs_lists,
+            prompt_logprobs_dict=prompt_logprobs_dict,
+            pooler_output=(
+                pooler_output
+                if self.vllm_config.model_config.engine_output_type != "text"
+                else None
+            ),
+            kv_connector_output=kv_connector_output,
+            num_nans_in_logits=num_nans_in_logits,
+        )
+
+        if not self.use_async_scheduling:
+            return output
+
+        return AsyncGPUModelRunnerOutput(
+            model_runner_output=output,
+            sampled_token_ids=sampler_output.sampled_token_ids,
+            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:
+            return
+        existing = getattr(req_state, "additional_information_cpu", {})
+        if not isinstance(existing, dict):
+            existing = {}
+        merged = dict(existing)
+        for k, v in upd.items():
+            if isinstance(v, torch.Tensor):
+                merged[k] = v.detach().to("cpu").contiguous()
+            elif isinstance(v, list):
+                new_list = []
+                for item in v:
+                    if isinstance(item, torch.Tensor):
+                        new_list.append(item.detach().to("cpu").contiguous())
+                    else:
+                        new_list.append(item)
+                merged[k] = new_list
+            else:
+                merged[k] = v
+        setattr(req_state, "additional_information_cpu", merged)

vllm_omni/worker/cpu/cpu_ar_worker.py

@@ -0,0 +1,63 @@
+import os
+import platform
+
+from vllm import envs
+import torch
+from vllm.logger import init_logger
+from vllm.model_executor import set_random_seed
+from vllm.platforms import current_platform, CpuArchEnum
+
+
+from vllm.v1.worker.cpu_worker import CPUWorker
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+
+from vllm_omni.worker.cpu.cpu_ar_model_runner import CPUARModelRunner
+
+logger = init_logger(__name__)
+
+
+class CPUARWorker(CPUWorker):
+    def init_device(self):
+        # Setup OpenMP threads affinity.
+        omp_cpuids = envs.VLLM_CPU_OMP_THREADS_BIND
+        if omp_cpuids == "auto" and platform.system() == "Linux":
+            cpu_arch = current_platform.get_cpu_architecture()
+            if cpu_arch in (CpuArchEnum.POWERPC, CpuArchEnum.S390X):
+                # For S390X/POWERPC SMT-8/4/2
+                self.local_omp_cpuid = self._get_autobind_cpu_ids(
+                    lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4])
+            elif current_platform.get_cpu_architecture() == CpuArchEnum.X86:
+                # For x86 SMT-2, use 1 CPU per core
+                self.local_omp_cpuid = self._get_autobind_cpu_ids(
+                    lambda cpus: cpus[-1:])
+            else:
+                self.local_omp_cpuid = "all"
+        else:
+            local_dp_rank = self.parallel_config.data_parallel_rank_local
+            omp_cpuids = omp_cpuids.split("|")
+            if local_dp_rank is not None:
+                world_size = self.parallel_config.world_size
+                omp_cpuids = omp_cpuids[local_dp_rank *
+                                        world_size:(local_dp_rank + 1) *
+                                                   world_size]
+            self.local_omp_cpuid = omp_cpuids[self.rank]
+
+        if self.local_omp_cpuid != "all":
+            ret = torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)
+            if ret:
+                logger.info(ret)
+
+        # Note: unique identifier for creating allreduce shared memory
+        os.environ["VLLM_DIST_IDENT"] = self.distributed_init_method.split(
+            ":")[-1]
+        # Initialize the distributed environment.
+        init_worker_distributed_environment(self.vllm_config, self.rank,
+                                            self.distributed_init_method,
+                                            self.local_rank,
+                                            current_platform.dist_backend)
+        # Set random seed.
+        set_random_seed(self.model_config.seed)
+
+        # Construct the model runner
+        self.model_runner: CPUARModelRunner = CPUARModelRunner(
+            self.vllm_config, torch.device("cpu"))

vllm_omni/worker/cpu/cpu_diffusion_model_runner.py

@@ -0,0 +1,620 @@
+from __future__ import annotations
+
+import gc
+import logging
+from typing import Any, List, Optional, Union
+
+import numpy as np
+import torch
+
+from vllm.config import CUDAGraphMode
+from vllm.forward_context import BatchDescriptor
+from vllm.multimodal.inputs import MultiModalKwargs
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.spec_decode.eagle import EagleProposer
+from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker.gpu_model_runner import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    AsyncGPUModelRunnerOutput,
+    IntermediateTensors,
+    PerLayerAttnMetadata,
+    get_pp_group,
+    set_forward_context,
+)
+from vllm.v1.worker.ubatch_utils import UBatchSlices
+from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.cpu.cpu_model_runner import OmniCPUModelRunner
+
+logger = logging.getLogger(__name__)
+
+
+class CPUDiffusionModelRunner(OmniCPUModelRunner):
+    """Diffusion model runner for vLLM-omni (non-autoregressive).
+
+    - Reuses CPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
+    - Does not compute logits or perform token sampling.
+    - Executes diffusion process and returns tensors via `pooler_output`.
+    """
+
+    def _preprocess(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        ubatch_slices: Optional[UBatchSlices] = None,
+        num_tokens_after_padding: Optional[torch.Tensor] = None,
+    ) -> tuple[
+        int,
+        int,
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        torch.Tensor,
+        Optional[IntermediateTensors],
+        dict[str, Any],
+    ]:
+
+        # Input token count for this iteration (not used by diffusion, but
+        # retained to keep DP padding/ordering consistent)
+        num_input_tokens = scheduler_output.total_num_scheduled_tokens
+        num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
+        num_input_tokens += num_pad
+
+        # _prepare_inputs may reorder the batch, so we must gather multi
+        # modal outputs after that to ensure the correct order
+        if (
+            self.supports_mm_inputs
+            and get_pp_group().is_first_rank
+            and not self.model_config.is_encoder_decoder
+        ):
+            # Run the multimodal encoder if any.
+            self._execute_mm_encoder(scheduler_output)
+            mm_embeds = self._gather_mm_embeddings(scheduler_output)
+
+            # NOTE(woosuk): To unify token ids and soft tokens (vision
+            # embeddings), we always use embeddings (rather than token ids)
+            # as input to the multimodal model, even when the input is text.
+            inputs_embeds_scheduled = self.model.get_input_embeddings(
+                input_ids=self.input_ids.gpu[:num_input_tokens],
+                multimodal_embeddings=mm_embeds or None,
+            )
+
+            # TODO(woosuk): Avoid the copy. Optimize.
+            self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled)
+
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = {
+                **self._init_model_kwargs(num_input_tokens),
+                **self._extract_mm_kwargs(scheduler_output),
+            }
+        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
+            # Get the input embeddings for the tokens that are not input embeds,
+            # then put them into the appropriate positions.
+            # TODO(qthequartermasterman): Since even when prompt embeds are
+            # enabled, (a) not all requests will use prompt embeds, and (b)
+            # after the initial prompt is processed, the rest of the generated
+            # tokens will be token ids, it is not desirable to have the
+            # embedding layer outside of the CUDA graph all the time. The v0
+            # engine avoids this by "double compiling" the CUDA graph, once
+            # with input_ids and again with inputs_embeds, for all num_tokens.
+            # If a batch only has token ids, then including the embedding layer
+            # in the CUDA graph will be more performant (like in the else case
+            # below).
+            token_ids_idx = (
+                self.is_token_ids.gpu[:num_input_tokens]
+                .nonzero(as_tuple=False)
+                .squeeze(1)
+            )
+            # Some tokens ids may need to become embeds
+            if token_ids_idx.numel() > 0:
+                token_ids = self.input_ids.gpu[token_ids_idx]
+                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
+                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
+
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+        else:
+            # For text-only models, we use token ids as input.
+            # While it is possible to use embeddings as input just like the
+            # multimodal models, it is not desirable for performance since
+            # then the embedding layer is not included in the CUDA graph.
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = None
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+        if self.uses_mrope:
+            positions = self.mrope_positions.gpu[:, :num_input_tokens]
+        else:
+            positions = self.positions.gpu[:num_input_tokens]
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True
+            )
+
+        if (
+            self.model_config.is_encoder_decoder
+            and scheduler_output.scheduled_encoder_inputs
+        ):
+            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
+            model_kwargs.update(encoder_inputs)
+
+        return (
+            num_input_tokens,
+            num_input_tokens,
+            num_tokens_after_padding,
+            input_ids,
+            inputs_embeds,
+            positions,
+            intermediate_tensors,
+            model_kwargs,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
+        with record_function_or_nullcontext("Preprocess"):
+            with self.synchronize_input_prep():
+                # Update persistent batch states.
+                self._update_states(scheduler_output)
+                if not scheduler_output.total_num_scheduled_tokens:
+                    return EMPTY_MODEL_RUNNER_OUTPUT
+
+                # Prepare the decoder inputs.
+                (
+                    attn_metadata,
+                    logits_indices,
+                    spec_decode_metadata,
+                    num_scheduled_tokens_np,
+                    spec_decode_common_attn_metadata,
+                    max_query_len,
+                    ubatch_slices,
+                    num_tokens_after_padding,
+                ) = self._prepare_inputs(scheduler_output)
+
+            # Input token count for this iteration (not used by diffusion, but
+            # retained to keep DP padding/ordering consistent)
+            (
+                num_scheduled_tokens,
+                num_input_tokens,
+                num_tokens_across_dp,
+                input_ids,
+                inputs_embeds,
+                positions,
+                intermediate_tensors,
+                model_kwargs,
+            ) = self._preprocess(
+                scheduler_output,
+                intermediate_tensors,
+                ubatch_slices,
+                num_tokens_after_padding,
+            )
+
+        cudagraph_runtime_mode = CUDAGraphMode.NONE
+        # Run the model.
+        # Use persistent buffers for CUDA graphs.
+        with (
+            set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
+                batch_descriptor=None,
+                ubatch_slices=ubatch_slices,
+            ),
+            record_function_or_nullcontext("Forward"),
+            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+        ):
+
+            outputs = self._run_diffusion(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+                multimodal_kwargs=model_kwargs,
+                logits_indices=logits_indices,
+            )
+
+        _, multimodal_outputs = self.extract_multimodal_outputs(outputs)
+        # Ensure one tensor per request, map to CPU for output struct
+        pooler_output: List[Optional[torch.Tensor]] = []
+        if isinstance(multimodal_outputs, torch.Tensor):
+            # If model returned a single stacked tensor, split by requests
+            assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
+            for i in range(self.input_batch.num_reqs):
+                pooler_output.append(
+                    multimodal_outputs[i].detach().to("cpu").contiguous()
+                )
+        elif isinstance(multimodal_outputs, list):
+            for out in multimodal_outputs:
+                pooler_output.append(
+                    out.detach().to("cpu").contiguous() if out is not None else None
+                )
+        elif isinstance(multimodal_outputs, dict):
+            for out in multimodal_outputs.values():
+                pooler_output.append(
+                    out.detach().to("cpu").contiguous() if out is not None else None
+                )
+        else:
+            raise RuntimeError("Unsupported diffusion output type")
+
+        output = OmniModelRunnerOutput(
+            req_ids=self.input_batch.req_ids,
+            req_id_to_index=self.input_batch.req_id_to_index,
+            sampled_token_ids=[],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=pooler_output,
+            kv_connector_output=kv_connector_output,
+            num_nans_in_logits={},
+        )
+
+        if not self.use_async_scheduling:
+            return output
+
+        return AsyncGPUModelRunnerOutput(
+            model_runner_output=output,
+            sampled_token_ids=[],
+            invalid_req_indices=[],
+            async_output_copy_stream=self.async_output_copy_stream,
+        )
+
+    def _run_diffusion(
+        self,
+        *,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors],
+        inputs_embeds: Optional[torch.Tensor],
+        multimodal_kwargs: dict,
+        logits_indices: torch.Tensor,
+    ) -> Union[torch.Tensor, list[torch.Tensor]]:
+        """Runs the diffusion process and returns per-request tensors.
+
+        Tries model interfaces in the following order for maximal compatibility:
+        1) model.sample(condition=..., **kwargs)
+        2) model.forward(condition=..., **kwargs)
+        3) model.diffuse(condition=..., **kwargs)
+        """
+        # Keep inputs identical to AR runner
+        kwargs = dict(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
+            sampling_metadata=self.input_batch.sampling_metadata,
+            logits_index=logits_indices,
+            sampler=self.sampler,
+        )
+
+        # For Qwen 2.5 Omni's current implementation, we only support the forward method
+        if hasattr(self.model, "forward"):
+            return self.model.forward(**kwargs)
+        # TODO: add the diffuse method for other models
+
+        raise RuntimeError(
+            "The loaded model does not expose diffusion interfaces 'sample', "
+            "'forward', or 'diffuse'. Please implement one of them or adapt the runner."
+        )
+
+    @torch.inference_mode()
+    def _dummy_run(
+        self,
+        num_tokens: int,
+        cudagraph_runtime_mode: Optional[CUDAGraphMode] = None,
+        force_attention: bool = False,
+        uniform_decode: bool = False,
+        skip_eplb: bool = False,
+        is_profile: bool = False,
+        create_mixed_batch: bool = False,
+        remove_lora: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Run a dummy forward pass to warm up/profile run or capture the
+        CUDA graph for the model.
+
+        Args:
+            num_tokens: Number of tokens to run the dummy forward pass.
+            cudagraph_runtime_mode: used to control the behavior.
+                - if not set will determine the cudagraph mode based on using
+                    the self.cudagraph_dispatcher.
+                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
+                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
+                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
+                    needed.
+            force_attention: If True, always create attention metadata. Used to
+                warm up attention backend when mode is NONE.
+            uniform_decode: If True, the batch is a uniform decode batch.
+            skip_eplb: If True, skip EPLB state update.
+            is_profile: If True, this is a profile run.
+            create_mixed_batch: If True, create a mixed batch with both decode
+                (1 token) and prefill (multiple tokens) requests.
+            remove_lora: If False, dummy LoRAs are not destroyed after the run
+        """
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
+            CUDAGraphMode.NONE,
+            CUDAGraphMode.PIECEWISE,
+            CUDAGraphMode.FULL,
+        }
+
+        # If cudagraph_mode.decode_mode() == FULL and
+        # cudagraph_mode.separate_routine(). This means that we are using
+        # different graphs and/or modes for mixed prefill-decode batches vs.
+        # uniform decode batches. A uniform decode batch means that all
+        # requests have identical query length, except a potential virtual
+        # request (shorter) in the batch account for padding.
+        # Uniform decode batch could either be common pure decode, where
+        # max_query_len == 1, or speculative decode, where
+        # max_query_len == 1 + num_spec_decode_tokens.
+
+        # When setting max_query_len = 1, we switch to and capture the optimized
+        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
+        # for GQA/MQA.
+        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
+
+        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
+        # for dummy run with LoRA so that the num_reqs collectively
+        # has num_tokens in total.
+        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+        max_num_reqs = self.scheduler_config.max_num_seqs
+
+        num_reqs = min(num_tokens, max_num_reqs)
+        min_tokens_per_req = num_tokens // num_reqs
+        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
+        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+
+        assert sum(num_scheduled_tokens_list) == num_tokens
+        assert len(num_scheduled_tokens_list) == num_reqs
+        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
+
+        num_tokens_after_padding = None
+
+        # If we failed to microbatch, currently need to resynchronize
+        # TODO(lucas,sage): we should be able to avoid this second sync by
+        #  refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
+        #  a single `coordinate_batch_across_dp` function.
+        if num_tokens_after_padding is None:
+            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
+            num_tokens_after_padding = num_tokens + num_pad
+        else:
+            num_tokens_across_dp = num_tokens_after_padding
+            num_tokens_after_padding = int(num_tokens_after_padding[0].item())
+
+        attn_metadata: Optional[PerLayerAttnMetadata] = None
+
+        # If force_attention is True, we always capture attention. Otherwise,
+        # it only happens for cudagraph_runtime_mode=FULL.
+        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
+            attn_metadata = {}
+
+            seq_lens = max_query_len
+            self.seq_lens.np[:num_reqs] = seq_lens
+            self.seq_lens.np[num_reqs:] = 0
+            self.seq_lens.copy_to_gpu()
+
+            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+            self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+            self.query_start_loc.copy_to_gpu()
+
+            for kv_cache_group_id, kv_cache_group_spec in enumerate(
+                self.kv_cache_config.kv_cache_groups
+            ):
+                common_attn_metadata = CommonAttentionMetadata(
+                    query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
+                    query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
+                    seq_lens=self.seq_lens.gpu[:num_reqs],
+                    seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
+                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
+                        :num_reqs
+                    ],
+                    num_reqs=num_reqs,
+                    num_actual_tokens=num_tokens,
+                    max_query_len=max_query_len,
+                    max_seq_len=self.max_model_len,
+                    block_table_tensor=self.input_batch.block_table[
+                        kv_cache_group_id
+                    ].get_device_tensor(num_reqs),
+                    slot_mapping=self.input_batch.block_table[
+                        kv_cache_group_id
+                    ].slot_mapping.gpu[:num_tokens],
+                    causal=True,
+                )
+                for attn_group in self.attn_groups[kv_cache_group_id]:
+
+                    assert type(attn_metadata) is dict
+                    attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(  # noqa: E501
+                        common_attn_metadata
+                    )
+                    for layer_name in attn_group.layer_names:
+                        attn_metadata[layer_name] = attn_metadata_i
+
+        with self.maybe_dummy_run_with_lora(
+            self.lora_config, num_scheduled_tokens, remove_lora
+        ):
+            model_kwargs = self._init_model_kwargs(num_tokens)
+            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                model_kwargs = {
+                    **model_kwargs,
+                    **self._dummy_mm_kwargs(num_reqs),
+                }
+            elif self.enable_prompt_embeds:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                model_kwargs = self._init_model_kwargs(num_tokens)
+            else:
+                input_ids = self.input_ids.gpu[:num_tokens]
+                inputs_embeds = None
+
+            if self.uses_mrope:
+                positions = self.mrope_positions.gpu[:, :num_tokens]
+            else:
+                positions = self.positions.gpu[:num_tokens]
+
+            if get_pp_group().is_first_rank:
+                intermediate_tensors = None
+            else:
+                if self.intermediate_tensors is None:
+                    self.intermediate_tensors = (
+                        self.model.make_empty_intermediate_tensors(
+                            batch_size=self.max_num_tokens,
+                            dtype=self.model_config.dtype,
+                            device=self.device,
+                        )
+                    )
+
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                    num_tokens, None, False
+                )
+
+            # filter out the valid batch descriptor
+            _cg_mode, batch_descriptor = (
+                self.cudagraph_dispatcher.dispatch(
+                    BatchDescriptor(
+                        num_tokens=num_tokens_after_padding,
+                        uniform_decode=uniform_decode,
+                    )
+                )
+                if not is_profile
+                else (CUDAGraphMode.NONE, None)
+            )
+            if cudagraph_runtime_mode is not None:
+                # we allow forcing NONE when the dispatcher disagrees to support
+                # warm ups for cudagraph capture
+                assert (
+                    cudagraph_runtime_mode == CUDAGraphMode.NONE
+                    or cudagraph_runtime_mode == _cg_mode
+                ), (
+                    f"Cudagraph runtime mode mismatch at dummy_run. "
+                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
+                )
+            else:
+                cudagraph_runtime_mode = _cg_mode
+
+            with self.maybe_randomize_inputs(input_ids), set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_tokens_after_padding,
+                num_tokens_across_dp=num_tokens_across_dp,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
+                batch_descriptor=batch_descriptor,
+            ):
+                outputs = self.model(
+                    input_ids=input_ids,
+                    positions=positions,
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                    **model_kwargs,
+                )
+
+            if self.use_aux_hidden_state_outputs:
+                hidden_states, _ = outputs
+            else:
+                hidden_states = outputs
+
+            if self.speculative_config and self.speculative_config.use_eagle():
+                assert isinstance(self.drafter, EagleProposer)
+                self.drafter.dummy_run(num_tokens)
+
+        # This is necessary to avoid blocking DP.
+        # For dummy runs, we typically skip EPLB since we don't have any real
+        # requests to process.
+        # However, in DP settings, there may be cases when some DP ranks do
+        # not have any requests to process, so they're executing dummy batches.
+        # In such cases, we still have to trigger EPLB to make sure
+        # ranks execute the rearrangement in synchronization.
+
+        # logit_indices = np.cumsum(num_scheduled_tokens) - 1
+        hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
+        return hidden_states, None
+
+    @torch.inference_mode()
+    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+        logger.warning("Dummy sampler run is not implemented for diffusion model")
+        return None
+
+    def profile_run(self) -> None:
+        # Profile with multimodal encoder & encoder cache.
+        if self.supports_mm_inputs:
+            if self.model_config.multimodal_config.skip_mm_profiling:
+                logger.info(
+                    "Skipping memory profiling for multimodal encoder and "
+                    "encoder cache."
+                )
+            else:
+                mm_budget = self.mm_budget
+                assert mm_budget is not None
+
+                if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
+                    # NOTE: Currently model is profiled with a single non-text
+                    # modality with the max possible input tokens even when
+                    # it supports multiple.
+                    dummy_modality = mm_budget.get_modality_with_max_tokens()
+                    max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[
+                        dummy_modality
+                    ]
+
+                    logger.info(
+                        "Encoder cache will be initialized with a budget of "
+                        "%s tokens, and profiled with %s %s items of the "
+                        "maximum feature size.",
+                        encoder_budget,
+                        max_mm_items_per_batch,
+                        dummy_modality,
+                    )
+
+                    # Create dummy batch of multimodal inputs.
+                    batched_dummy_mm_inputs = self._get_mm_dummy_batch(
+                        dummy_modality,
+                        max_mm_items_per_batch,
+                    )
+
+                    # Run multimodal encoder.
+                    dummy_encoder_outputs = self.model.get_multimodal_embeddings(
+                        **batched_dummy_mm_inputs
+                    )
+
+                    sanity_check_mm_encoder_outputs(
+                        dummy_encoder_outputs,
+                        expected_num_items=max_mm_items_per_batch,
+                    )
+
+                    # NOTE: This happens when encoder cache needs to store
+                    # the embeddings that encoder outputs are scattered onto.
+                    # In this case we create dummy embeddings of size
+                    # (encode_budget, hidden_size) and scatter encoder
+                    # output into it.
+                    encoder_output_shape = dummy_encoder_outputs[0].shape
+                    if encoder_output_shape[0] < encoder_budget:
+                        expanded_outputs = []
+                        for output in dummy_encoder_outputs:
+                            expanded = output.new_zeros(
+                                (encoder_budget, encoder_output_shape[-1])
+                            )
+                            num_tokens = output.shape[0]
+                            expanded[:num_tokens].copy_(output)
+                            expanded_outputs.append(expanded)
+
+                        dummy_encoder_outputs = expanded_outputs
+
+                    # Cache the dummy encoder outputs.
+                    self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
+
+        # Add `is_profile` here to pre-allocate communication buffers
+        hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
+        if get_pp_group().is_last_rank:
+            pass
+        self._sync_device()
+        del hidden_states
+        self.encoder_cache.clear()
+        gc.collect()

vllm_omni/worker/cpu/cpu_diffusion_worker.py

@@ -0,0 +1,66 @@
+import gc
+import os
+import platform
+
+import torch
+from vllm import envs
+from vllm.logger import init_logger
+
+from vllm.model_executor import set_random_seed
+from vllm.platforms import current_platform, CpuArchEnum
+
+from vllm.v1.worker.cpu_worker import CPUWorker
+
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+
+from vllm_omni.worker.cpu.cpu_diffusion_model_runner import CPUDiffusionModelRunner
+
+
+logger = init_logger(__name__)
+
+
+class CPUDiffusionWorker(CPUWorker):
+    def init_device(self):
+        # Setup OpenMP threads affinity.
+        omp_cpuids = envs.VLLM_CPU_OMP_THREADS_BIND
+        if omp_cpuids == "auto" and platform.system() == "Linux":
+            cpu_arch = current_platform.get_cpu_architecture()
+            if cpu_arch in (CpuArchEnum.POWERPC, CpuArchEnum.S390X):
+                # For S390X/POWERPC SMT-8/4/2
+                self.local_omp_cpuid = self._get_autobind_cpu_ids(
+                    lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4])
+            elif current_platform.get_cpu_architecture() == CpuArchEnum.X86:
+                # For x86 SMT-2, use 1 CPU per core
+                self.local_omp_cpuid = self._get_autobind_cpu_ids(
+                    lambda cpus: cpus[-1:])
+            else:
+                self.local_omp_cpuid = "all"
+        else:
+            local_dp_rank = self.parallel_config.data_parallel_rank_local
+            omp_cpuids = omp_cpuids.split("|")
+            if local_dp_rank is not None:
+                world_size = self.parallel_config.world_size
+                omp_cpuids = omp_cpuids[local_dp_rank *
+                                        world_size:(local_dp_rank + 1) *
+                                                   world_size]
+            self.local_omp_cpuid = omp_cpuids[self.rank]
+
+        if self.local_omp_cpuid != "all":
+            ret = torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)
+            if ret:
+                logger.info(ret)
+
+        # Note: unique identifier for creating allreduce shared memory
+        os.environ["VLLM_DIST_IDENT"] = self.distributed_init_method.split(
+            ":")[-1]
+        # Initialize the distributed environment.
+        init_worker_distributed_environment(self.vllm_config, self.rank,
+                                            self.distributed_init_method,
+                                            self.local_rank,
+                                            current_platform.dist_backend)
+        # Set random seed.
+        set_random_seed(self.model_config.seed)
+
+        # Construct the model runner
+        self.model_runner: CPUDiffusionModelRunner = CPUDiffusionModelRunner(
+            self.vllm_config, torch.device("cpu"))

vllm_omni/worker/cpu/cpu_model_runner.py

@@ -0,0 +1,593 @@
+from typing import TYPE_CHECKING, List, Optional, Union, cast
+
+import numpy as np
+import torch
+
+from vllm.config import CUDAGraphMode
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.forward_context import BatchDescriptor, set_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces_base import VllmModelForPooling
+from vllm.sampling_params import SamplingType
+from vllm.utils import LazyLoader, cdiv
+from vllm.v1.attention.backends.utils import (
+    CommonAttentionMetadata,
+    split_attn_metadata,
+)
+from vllm.v1.spec_decode.eagle import EagleProposer
+# from vllm.v1.worker.cpu import CachedRequestState
+from vllm.v1.worker.cpu_model_runner import CPUModelRunner
+from vllm.v1.worker.gpu_input_batch import CachedRequestState
+from vllm.v1.worker.gpu_model_runner import PerLayerAttnMetadata
+from vllm.v1.worker.ubatch_splitting import ubatch_split
+
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.output import SchedulerOutput
+else:
+    xgr = LazyLoader("xgr", globals(), "xgrammar")
+    xgr_torch_compile = LazyLoader(
+        "xgr_torch_compile",
+        globals(),
+        "xgrammar.kernels.apply_token_bitmask_inplace_torch_compile",
+    )
+
+logger = init_logger(__name__)
+
+
+class OmniCPUModelRunner(CPUModelRunner):
+    def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
+        """Update the cached states and the persistent batch with the scheduler
+        output.
+
+        The updated states are used by the `_prepare_inputs` function to create
+        the input CPU tensors for the model.
+
+        The SamplingMetadata is updated and copied to the CPU if there is a
+        new/resumed/paused/finished request in the batch.
+        """
+        # Remove finished requests from the cached states.
+        for req_id in scheduler_output.finished_req_ids:
+            self.requests.pop(req_id, None)
+        # Remove the finished requests from the persistent batch.
+        # NOTE(woosuk): There could be an edge case where finished_req_ids and
+        # scheduled_req_ids overlap. This happens when a request is aborted and
+        # then resubmitted with the same ID. In this case, we treat them as two
+        # distinct requests - clearing the cached states for the first request
+        # and handling the second as a new request.
+        for req_id in scheduler_output.finished_req_ids:
+            self.input_batch.remove_request(req_id)
+
+        # Free the cached encoder outputs.
+        for mm_hash in scheduler_output.free_encoder_mm_hashes:
+            self.encoder_cache.pop(mm_hash, None)
+
+        # Remove the unscheduled requests from the persistent batch.
+        # NOTE(woosuk): The unscheduled requests are either preempted requests
+        # or running requests that are not scheduled in this step. We remove
+        # them from the persistent batch but keep their cached states since
+        # they will be scheduled again sometime in the future.
+        scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
+        cached_req_ids = self.input_batch.req_id_to_index.keys()
+        unscheduled_req_ids = cached_req_ids - scheduled_req_ids
+        # NOTE(woosuk): The persistent batch optimization assumes that
+        # consecutive batches contain mostly the same requests. If batches
+        # have low request overlap (e.g., alternating between two distinct
+        # sets of requests), this optimization becomes very inefficient.
+        for req_id in unscheduled_req_ids:
+            self.input_batch.remove_request(req_id)
+
+        reqs_to_add: list[CachedRequestState] = []
+        # Add new requests to the cached states.
+        for new_req_data in scheduler_output.scheduled_new_reqs:
+            req_id = new_req_data.req_id
+            sampling_params = new_req_data.sampling_params
+            pooling_params = new_req_data.pooling_params
+
+            if (
+                    sampling_params
+                    and sampling_params.sampling_type == SamplingType.RANDOM_SEED
+            ):
+                generator = torch.Generator(device=self.device)
+                generator.manual_seed(sampling_params.seed)
+            else:
+                generator = None
+
+            if self.is_pooling_model:
+                assert pooling_params is not None
+                task = pooling_params.task
+                assert task is not None, "You did not set `task` in the API"
+
+                model = cast(VllmModelForPooling, self.get_model())
+                to_update = model.pooler.get_pooling_updates(task)
+                to_update.apply(pooling_params)
+
+            req_state = CachedRequestState(
+                req_id=req_id,
+                prompt_token_ids=new_req_data.prompt_token_ids,
+                prompt_embeds=new_req_data.prompt_embeds,
+                mm_features=new_req_data.mm_features,
+                sampling_params=sampling_params,
+                pooling_params=pooling_params,
+                generator=generator,
+                block_ids=new_req_data.block_ids,
+                num_computed_tokens=new_req_data.num_computed_tokens,
+                output_token_ids=[],
+                lora_request=new_req_data.lora_request,
+            )
+            self.requests[req_id] = req_state
+
+            # If prompt embeddings are provided, decode and attach to inter_data
+            try:
+                if getattr(new_req_data, "prompt_embeds", None) is not None:
+                    payload = new_req_data.prompt_embeds
+                    dtype = getattr(np, payload.dtype)
+                    arr = np.frombuffer(payload.data, dtype=dtype)
+                    arr = arr.reshape(payload.shape)
+                    pe_cpu = torch.from_numpy(arr)
+                    # Store temporarily on CPU; later moved to device in builder
+                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
+                    # Also replace payload with Tensor for user visibility in
+                    # scheduler_output
+                    try:
+                        new_req_data.prompt_embeds = pe_cpu  # type: ignore[assignment]
+                    except Exception:
+                        pass
+            except Exception as e:
+                logger.error(f"Error decoding prompt embeds: {e}")
+            # Decode additional_information payloads (dictionary)
+            try:
+                if getattr(new_req_data, "additional_information", None) is not None:
+                    payload_info = new_req_data.additional_information
+                    info_dict = {}
+                    if isinstance(payload_info, dict):
+                        info_dict = payload_info
+                    else:
+                        from vllm_omni.engine import AdditionalInformationPayload
+
+                        if isinstance(payload_info, AdditionalInformationPayload):
+                            for k, entry in payload_info.entries.items():
+                                if entry.tensor_data is not None:
+                                    dt = np.dtype(
+                                        getattr(entry, "tensor_dtype", "float32")
+                                    )
+                                    arr = np.frombuffer(entry.tensor_data, dtype=dt)
+                                    arr = arr.reshape(entry.tensor_shape)
+                                    info_dict[k] = torch.from_numpy(arr)
+                                else:
+                                    info_dict[k] = entry.list_data
+                    if info_dict:
+                        setattr(
+                            self.requests[req_id],
+                            "additional_information_cpu",
+                            info_dict,
+                        )
+            except Exception as e:
+                logger.error(f"Error decoding additional information: {e}")
+                pass
+
+            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
+            if self.uses_mrope:
+                self._init_mrope_positions(req_state)
+
+            reqs_to_add.append(self.requests[req_id])
+
+        # Update the states of the running/resumed requests.
+        is_last_rank = get_pp_group().is_last_rank
+        req_data = scheduler_output.scheduled_cached_reqs
+        for i, req_id in enumerate(req_data.req_ids):
+            req_state = self.requests[req_id]
+            num_computed_tokens = req_data.num_computed_tokens[i]
+            new_block_ids = req_data.new_block_ids[i]
+            resumed_from_preemption = req_data.resumed_from_preemption[i]
+
+            # Update the cached states.
+            req_state.num_computed_tokens = num_computed_tokens
+
+            if not is_last_rank:
+                # When using PP, the scheduler sends the sampled tokens back,
+                # because there's no direct communication between the first-
+                # stage worker and the last-stage worker.
+                new_token_ids = req_data.new_token_ids[i]
+                # Add the sampled token(s) from the previous step (if any).
+                # This doesn't include "unverified" tokens like spec tokens.
+                num_new_tokens = (
+                        num_computed_tokens + len(new_token_ids) - req_state.num_tokens
+                )
+                if num_new_tokens == 1:
+                    # Avoid slicing list in most common case.
+                    req_state.output_token_ids.append(new_token_ids[-1])
+                elif num_new_tokens > 0:
+                    req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
+
+            # Update the block IDs.
+            if not resumed_from_preemption:
+                if new_block_ids is not None:
+                    # Append the new blocks to the existing block IDs.
+                    for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
+                        block_ids.extend(new_ids)
+            else:
+                assert new_block_ids is not None
+                # The request is resumed from preemption.
+                # Replace the existing block IDs with the new ones.
+                req_state.block_ids = new_block_ids
+
+            req_index = self.input_batch.req_id_to_index.get(req_id)
+            if req_index is None:
+                # The request is not in the persistent batch.
+                # The request was either preempted and resumed later, or was not
+                # scheduled in the previous step and needs to be added again.
+                reqs_to_add.append(req_state)
+                continue
+
+            # Update the persistent batch.
+            self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens
+            if new_block_ids is not None:
+                self.input_batch.block_table.append_row(new_block_ids, req_index)
+
+            # For the last rank, we don't need to update the token_ids_cpu
+            # because the sampled tokens are already cached.
+            if not is_last_rank:
+                # Add new_token_ids to token_ids_cpu.
+                start_token_index = num_computed_tokens
+                end_token_index = num_computed_tokens + len(new_token_ids)
+                self.input_batch.token_ids_cpu[
+                    req_index, start_token_index:end_token_index
+                ] = new_token_ids
+                self.input_batch.num_tokens_no_spec[req_index] = end_token_index
+                self.input_batch.num_tokens[req_index] = end_token_index
+
+            # Add spec_token_ids to token_ids_cpu.
+            spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
+                req_id, ()
+            )
+            if spec_token_ids:
+                num_spec_tokens = len(spec_token_ids)
+                start_index = self.input_batch.num_tokens_no_spec[req_index]
+                end_token_index = start_index + num_spec_tokens
+                self.input_batch.token_ids_cpu[
+                    req_index, start_index:end_token_index
+                ] = spec_token_ids
+                # NOTE(woosuk): `num_tokens` here may include spec tokens.
+                self.input_batch.num_tokens[req_index] += num_spec_tokens
+
+        # Add the new or resumed requests to the persistent batch.
+        # The smaller empty indices are filled first.
+        for request in reqs_to_add:
+            self.input_batch.add_request(request)
+
+        # Condense the batched states if there are gaps left by removed requests
+        self.input_batch.condense()
+        # Allow attention backend to reorder the batch, potentially
+        self._may_reorder_batch(scheduler_output)
+        # Refresh batch metadata with any pending updates.
+        self.input_batch.refresh_metadata()
+
+    @torch.inference_mode()
+    def extract_multimodal_outputs(
+            self, hidden_states: Union[torch.Tensor, List[torch.Tensor]]
+    ) -> dict:
+        if (
+                hasattr(self.model, "have_multimodal_outputs")
+                and self.model.have_multimodal_outputs
+        ):
+            text_hidden_states = hidden_states.text_hidden_states
+            multimodal_outputs = hidden_states.multimodal_outputs
+
+        elif isinstance(hidden_states, torch.Tensor):
+            text_hidden_states = hidden_states
+            multimodal_outputs = {}
+        elif isinstance(hidden_states, List):
+            text_hidden_states = hidden_states[0]
+            multimodal_outputs = {}
+        else:
+            raise ValueError(f"Invalid hidden states type: {type(hidden_states)}")
+        return text_hidden_states, multimodal_outputs
+
+    @torch.inference_mode()
+    def _dummy_run(
+            self,
+            num_tokens: int,
+            cudagraph_runtime_mode: Optional[CUDAGraphMode] = None,
+            force_attention: bool = False,
+            uniform_decode: bool = False,
+            allow_microbatching: bool = True,
+            skip_eplb: bool = False,
+            is_profile: bool = False,
+            create_mixed_batch: bool = False,
+            remove_lora: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Run a dummy forward pass to warm up/profile run or capture the
+        CUDA graph for the model.
+
+        Args:
+            num_tokens: Number of tokens to run the dummy forward pass.
+            cudagraph_runtime_mode: used to control the behavior.
+                - if not set will determine the cudagraph mode based on using
+                    the self.cudagraph_dispatcher.
+                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
+                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
+                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
+                    needed.
+            force_attention: If True, always create attention metadata. Used to
+                warm up attention backend when mode is NONE.
+            uniform_decode: If True, the batch is a uniform decode batch.
+            skip_eplb: If True, skip EPLB state update.
+            is_profile: If True, this is a profile run.
+            create_mixed_batch: If True, create a mixed batch with both decode
+                (1 token) and prefill (multiple tokens) requests.
+            remove_lora: If False, dummy LoRAs are not destroyed after the run
+        """
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
+            CUDAGraphMode.NONE,
+            CUDAGraphMode.PIECEWISE,
+            CUDAGraphMode.FULL,
+        }
+
+        # If cudagraph_mode.decode_mode() == FULL and
+        # cudagraph_mode.separate_routine(). This means that we are using
+        # different graphs and/or modes for mixed prefill-decode batches vs.
+        # uniform decode batches. A uniform decode batch means that all
+        # requests have identical query length, except a potential virtual
+        # request (shorter) in the batch account for padding.
+        # Uniform decode batch could either be common pure decode, where
+        # max_query_len == 1, or speculative decode, where
+        # max_query_len == 1 + num_spec_decode_tokens.
+
+        # When setting max_query_len = 1, we switch to and capture the optimized
+        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
+        # for GQA/MQA.
+        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
+
+        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
+        # for dummy run with LoRA so that the num_reqs collectively
+        # has num_tokens in total.
+        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+        max_num_reqs = self.scheduler_config.max_num_seqs
+        if create_mixed_batch:
+            assert not uniform_decode
+            # Create mixed batch:
+            # first half decode tokens, second half one prefill
+            num_decode_tokens = num_tokens // 2
+            num_prefill_tokens = num_tokens - num_decode_tokens
+            num_reqs = num_decode_tokens + 1
+
+            # Create decode requests (1 token each) followed by prefill request
+            num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens]
+            # Note: Overriding max_query_len to be the prefill tokens
+            max_query_len = num_prefill_tokens
+        elif uniform_decode:
+            assert not create_mixed_batch
+            num_reqs = cdiv(num_tokens, max_query_len)
+            num_scheduled_tokens_list = [max_query_len] * num_reqs
+            if num_tokens % max_query_len != 0:
+                num_scheduled_tokens_list[-1] = num_tokens % max_query_len
+        else:
+            num_reqs = min(num_tokens, max_num_reqs)
+            min_tokens_per_req = num_tokens // num_reqs
+            num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
+            num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+
+        assert sum(num_scheduled_tokens_list) == num_tokens
+        assert len(num_scheduled_tokens_list) == num_reqs
+        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
+
+        total_num_scheduled_tokens = int(num_scheduled_tokens.sum())
+
+        ubatch_slices = None
+        num_tokens_after_padding = None
+
+        # We currently only microbatch if the number of tokens is
+        # over a certain threshold.
+        if self.parallel_config.enable_dbo and allow_microbatching:
+            ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split(
+                num_scheduled_tokens,
+                total_num_scheduled_tokens,
+                total_num_scheduled_tokens,
+                uniform_decode=uniform_decode,
+                vllm_config=self.vllm_config,
+            )
+            # Currently when DBO is enabled `ubatch_split` returns
+            # the num_tokens_after_padding for a single ubatch, but we have 2
+            # TODO(sage,lucas): this is cruft that should be addressed in the
+            # padding refactor.
+            if ubatch_num_tokens_after_padding is not None:
+                num_tokens_after_padding = ubatch_num_tokens_after_padding * 2
+
+        # If we failed to microbatch, currently need to resynchronize
+        # TODO(lucas,sage): we should be able to avoid this second sync by
+        #  refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
+        #  a single `coordinate_batch_across_dp` function.
+        if num_tokens_after_padding is None:
+            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
+            num_tokens_after_padding = num_tokens + num_pad
+        else:
+            num_tokens_across_dp = num_tokens_after_padding
+            num_tokens_after_padding = int(num_tokens_after_padding[0].item())
+
+        attn_metadata: Optional[PerLayerAttnMetadata] = None
+
+        # If force_attention is True, we always capture attention. Otherwise,
+        # it only happens for cudagraph_runtime_mode=FULL.
+        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
+            attn_metadata = {}
+            if ubatch_slices is not None:
+                attn_metadata = [dict() for _ in range(len(ubatch_slices))]
+
+            if create_mixed_batch:
+                # In the mixed batch mode (used for FI warmup), we use
+                # shorter sequence lengths to run faster.
+                # TODO(luka) better system for describing dummy batches
+                seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
+            else:
+                seq_lens = max_query_len
+            self.seq_lens.np[:num_reqs] = seq_lens
+            self.seq_lens.np[num_reqs:] = 0
+            self.seq_lens.copy_to_gpu()
+
+            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+            self.query_start_loc.np[1: num_reqs + 1] = cum_num_tokens
+            self.query_start_loc.copy_to_gpu()
+
+            for kv_cache_group_id, kv_cache_group_spec in enumerate(
+                    self.kv_cache_config.kv_cache_groups
+            ):
+                common_attn_metadata = CommonAttentionMetadata(
+                    query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
+                    query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
+                    seq_lens=self.seq_lens.gpu[:num_reqs],
+                    seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
+                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
+                        :num_reqs
+                    ],
+                    num_reqs=num_reqs,
+                    num_actual_tokens=num_tokens,
+                    max_query_len=max_query_len,
+                    max_seq_len=self.max_model_len,
+                    block_table_tensor=self.input_batch.block_table[
+                        kv_cache_group_id
+                    ].get_device_tensor(num_reqs),
+                    slot_mapping=self.input_batch.block_table[
+                        kv_cache_group_id
+                    ].slot_mapping.gpu[:num_tokens],
+                    causal=True,
+                )
+                for attn_group in self.attn_groups[kv_cache_group_id]:
+                    if ubatch_slices is not None:
+                        common_attn_metadata_list = split_attn_metadata(
+                            ubatch_slices, common_attn_metadata
+                        )
+                        for ubid, common_attn_metadata in enumerate(
+                                common_attn_metadata_list
+                        ):
+                            assert common_attn_metadata.max_query_len == 1
+                            attn_metadata_i = attn_group.get_metadata_builder(
+                                ubatch_id=ubid
+                            ).build_for_cudagraph_capture(common_attn_metadata)
+                            for layer_name in attn_group.layer_names:
+                                assert type(attn_metadata) is list
+                                attn_metadata[ubid][layer_name] = attn_metadata_i
+                    else:
+                        assert type(attn_metadata) is dict
+                        attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(  # noqa: E501
+                            common_attn_metadata
+                        )
+                        for layer_name in attn_group.layer_names:
+                            attn_metadata[layer_name] = attn_metadata_i
+
+        with self.maybe_dummy_run_with_lora(
+                self.lora_config, num_scheduled_tokens, remove_lora
+        ):
+            model_kwargs = self._init_model_kwargs(num_tokens)
+            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                model_kwargs = {
+                    **model_kwargs,
+                    **self._dummy_mm_kwargs(num_reqs),
+                }
+            elif self.enable_prompt_embeds:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                model_kwargs = self._init_model_kwargs(num_tokens)
+            else:
+                input_ids = self.input_ids.gpu[:num_tokens]
+                inputs_embeds = None
+
+            if self.uses_mrope:
+                positions = self.mrope_positions.gpu[:, :num_tokens]
+            else:
+                positions = self.positions.gpu[:num_tokens]
+
+            if get_pp_group().is_first_rank:
+                intermediate_tensors = None
+            else:
+                if self.intermediate_tensors is None:
+                    self.intermediate_tensors = (
+                        self.model.make_empty_intermediate_tensors(
+                            batch_size=self.max_num_tokens,
+                            dtype=self.model_config.dtype,
+                            device=self.device,
+                        )
+                    )
+
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                    num_tokens, None, False
+                )
+
+            # filter out the valid batch descriptor
+            _cg_mode, batch_descriptor = (
+                self.cudagraph_dispatcher.dispatch(
+                    BatchDescriptor(
+                        num_tokens=num_tokens_after_padding,
+                        uniform_decode=uniform_decode,
+                    )
+                )
+                if not is_profile
+                else (CUDAGraphMode.NONE, None)
+            )
+            if cudagraph_runtime_mode is not None:
+                # we allow forcing NONE when the dispatcher disagrees to support
+                # warm ups for cudagraph capture
+                assert (
+                        cudagraph_runtime_mode == CUDAGraphMode.NONE
+                        or cudagraph_runtime_mode == _cg_mode
+                ), (
+                    f"Cudagraph runtime mode mismatch at dummy_run. "
+                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
+                )
+            else:
+                cudagraph_runtime_mode = _cg_mode
+
+            if ubatch_slices is not None:
+                # Adjust values to reflect a single ubatch.
+                # TODO(sage,lucas): this is cruft that should be addressed in
+                #  the padding refactor.
+                num_tokens_after_padding = ubatch_slices[0].num_tokens
+                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,
+            ):
+                outputs = self.model(
+                    input_ids=input_ids,
+                    positions=positions,
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                    **model_kwargs,
+                )
+
+            if self.use_aux_hidden_state_outputs:
+                hidden_states, _ = outputs
+            else:
+                hidden_states = outputs
+
+            if self.speculative_config and self.speculative_config.use_eagle():
+                assert isinstance(self.drafter, EagleProposer)
+                self.drafter.dummy_run(num_tokens)
+
+        # This is necessary to avoid blocking DP.
+        # For dummy runs, we typically skip EPLB since we don't have any real
+        # requests to process.
+        # However, in DP settings, there may be cases when some DP ranks do
+        # not have any requests to process, so they're executing dummy batches.
+        # In such cases, we still have to trigger EPLB to make sure
+        # ranks execute the rearrangement in synchronization.
+        if not skip_eplb:
+            self.eplb_step(is_dummy=True, is_profile=is_profile)
+
+        logit_indices = np.cumsum(num_scheduled_tokens) - 1
+        hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+            hidden_states
+        )
+        if len(hidden_states.shape) == 3:
+            return hidden_states[0], hidden_states[0][logit_indices]
+        else:
+            return hidden_states, hidden_states[logit_indices]

vllm_omni/worker/gpu/gpu_ar_model_runner.py

@@ -0,0 +1,649 @@
+"""AR GPU Model Runner for vLLM-omni.
+
+Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
+and also outputs sampled tokens.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Union
+
+import numpy as np
+import torch
+
+from vllm.forward_context import BatchDescriptor
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import AsyncModelRunnerOutput
+from vllm.v1.structured_output.utils import apply_grammar_bitmask
+from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker.gpu_model_runner import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    AsyncGPUModelRunnerOutput,
+    IntermediateTensors,
+    get_pp_group,
+    get_tp_group,
+    has_kv_transfer_group,
+    set_forward_context,
+)
+from vllm.v1.worker.ubatch_utils import UBatchSlices
+from vllm.v1.worker.utils import is_residual_scattered_for_sp
+from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.gpu.gpu_model_runner import OmniGPUModelRunner
+
+logger = init_logger(__name__)
+
+
+class GPUARModelRunner(OmniGPUModelRunner):
+    """Autoregressive GPU model runner that returns hidden states per request.
+
+    This runner follows the same preparation and forward path as GPUModelRunner
+    (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs),
+    and additionally performs lightweight sampling so that sampled tokens are
+    available in outputs. Hidden representations are taken at the same indices
+    that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`).
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
+
+    def _preprocess(
+        self,
+        scheduler_output: "SchedulerOutput",
+        num_scheduled_tokens_np: np.ndarray,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        ubatch_slices: Optional[UBatchSlices] = None,
+        num_tokens_after_padding: Optional[torch.Tensor] = None,
+    ) -> tuple[
+        int,
+        int,
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        torch.Tensor,
+        Optional[IntermediateTensors],
+        dict[str, Any],
+        Optional[dict[str, dict]],
+    ]:
+
+        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        if ubatch_slices:
+            assert num_tokens_after_padding is not None
+            num_input_tokens = int(num_tokens_after_padding[0].item() * 2)
+            self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
+        elif ubatch_slices is None:
+            num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)
+            num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
+            num_input_tokens += num_pad
+
+        # _prepare_inputs may reorder the batch, so we must gather multi
+        # modal outputs after that to ensure the correct order
+        per_req_additional_information: Optional[dict[str, dict]] = None
+        if (
+            self.supports_mm_inputs
+            and get_pp_group().is_first_rank
+            and not self.model_config.is_encoder_decoder
+        ):
+            # Run the multimodal encoder if any.
+            self._execute_mm_encoder(scheduler_output)
+            mm_embeds = self._gather_mm_embeddings(scheduler_output)
+
+            # NOTE(woosuk): To unify token ids and soft tokens (vision
+            # embeddings), we always use embeddings (rather than token ids)
+            # as input to the multimodal model, even when the input is text.
+            inputs_embeds_scheduled = self.model.get_input_embeddings(
+                input_ids=self.input_ids.gpu[:num_scheduled_tokens],
+                multimodal_embeddings=mm_embeds or None,
+            )
+
+            # TODO(woosuk): Avoid the copy. Optimize.
+            self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
+
+            # Reset per-step additional information collector (deprecated concat path)
+            if hasattr(self, "_forward_additional_information"):
+                self._forward_additional_information = None
+            # New: per-request additional information for this step
+            per_req_additional_information = {}
+
+            # Overlay custom prompt_embeds per request for the prompt portion;
+            # collect additional_information (tensor/list) for prefill portion only
+            for req_index, req_id in enumerate(self.input_batch.req_ids):
+                req_state = self.requests[req_id]
+                pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
+                addi_cpu = getattr(req_state, "additional_information_cpu", None)
+                num_computed_tokens = int(
+                    self.input_batch.num_computed_tokens_cpu[req_index]
+                )
+                prompt_len = len(req_state.prompt_token_ids)
+                prompt_remaining = max(0, prompt_len - num_computed_tokens)
+                sched_tokens = int(num_scheduled_tokens_np[req_index])
+                overlay_len = min(sched_tokens, prompt_remaining)
+                if overlay_len <= 0:
+                    continue
+                if pe_cpu is not None:
+                    src = pe_cpu[
+                        num_computed_tokens : num_computed_tokens + overlay_len
+                    ].to(dtype=self.dtype, device=self.device, non_blocking=True)
+                    start_offset = int(self.query_start_loc.cpu[req_index])
+                    self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(
+                        src
+                    )
+                # Build per-request additional information (no cross-request concat)
+                if addi_cpu is not None and isinstance(addi_cpu, dict):
+                    req_info: dict[str, object] = {}
+                    for k, v in addi_cpu.items():
+                        if isinstance(v, torch.Tensor):
+                            # 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()
+                                except Exception:
+                                    seg = v.detach().to("cpu").contiguous()
+                                req_info[k] = seg
+                            else:
+                                req_info[k] = v.detach().to("cpu").contiguous()
+                        elif isinstance(v, list):
+                            req_info[k] = v
+                        else:
+                            req_info[k] = v
+                    per_req_additional_information[req_id] = req_info
+
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = {
+                **self._init_model_kwargs(num_scheduled_tokens),
+                **self._extract_mm_kwargs(scheduler_output),
+            }
+        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
+            # Get the input embeddings for the tokens that are not input embeds,
+            # then put them into the appropriate positions.
+            # TODO(qthequartermasterman): Since even when prompt embeds are
+            # enabled, (a) not all requests will use prompt embeds, and (b)
+            # after the initial prompt is processed, the rest of the generated
+            # tokens will be token ids, it is not desirable to have the
+            # embedding layer outside of the CUDA graph all the time. The v0
+            # engine avoids this by "double compiling" the CUDA graph, once
+            # with input_ids and again with inputs_embeds, for all num_tokens.
+            # If a batch only has token ids, then including the embedding layer
+            # in the CUDA graph will be more performant (like in the else case
+            # below).
+            token_ids_idx = (
+                self.is_token_ids.gpu[:num_scheduled_tokens]
+                .nonzero(as_tuple=False)
+                .squeeze(1)
+            )
+            # Some tokens ids may need to become embeds
+            if token_ids_idx.numel() > 0:
+                token_ids = self.input_ids.gpu[token_ids_idx]
+                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
+                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
+
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+            input_ids = None
+        else:
+            # For text-only models, we use token ids as input.
+            # While it is possible to use embeddings as input just like the
+            # multimodal models, it is not desirable for performance since
+            # then the embedding layer is not included in the CUDA graph.
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = None
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+        if self.uses_mrope:
+            positions = self.mrope_positions.gpu[:, :num_input_tokens]
+        else:
+            positions = self.positions.gpu[:num_input_tokens]
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True
+            )
+
+        if (
+            self.model_config.is_encoder_decoder
+            and scheduler_output.scheduled_encoder_inputs
+        ):
+            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
+            model_kwargs.update(encoder_inputs)
+
+        return (
+            num_scheduled_tokens,
+            num_input_tokens,
+            num_tokens_after_padding,
+            input_ids,
+            inputs_embeds,
+            positions,
+            intermediate_tensors,
+            model_kwargs,
+            per_req_additional_information,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[OmniModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:
+        with record_function_or_nullcontext("Preprocess"):
+            with self.synchronize_input_prep():
+                # Update persistent batch states.
+                self._update_states(scheduler_output)
+
+                # Decode per-request prompt_embeds / additional_hidden_states payloads
+                # (if present) into CPU tensors
+                try:
+                    new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
+                    if new_reqs:
+                        for nr in new_reqs:
+                            req_id = getattr(nr, "req_id", None) or getattr(
+                                nr, "request_id", None
+                            )
+                            if req_id is None:
+                                continue
+                            # prompt_embeds
+                            payload_pe = getattr(nr, "prompt_embeds", None)
+                            if payload_pe is not None:
+                                if isinstance(payload_pe, torch.Tensor):
+                                    pe_cpu = payload_pe.detach().to("cpu").contiguous()
+                                elif isinstance(payload_pe, PromptEmbedsPayload):
+                                    dt = np.dtype(
+                                        getattr(payload_pe, "dtype", "float32")
+                                    )
+                                    arr = np.frombuffer(payload_pe.data, dtype=dt)
+                                    arr = arr.reshape(payload_pe.shape)
+                                    pe_cpu = torch.from_numpy(arr)
+                                else:
+                                    pe_cpu = None
+                                if pe_cpu is not None and req_id in self.requests:
+                                    setattr(
+                                        self.requests[req_id],
+                                        "prompt_embeds_cpu",
+                                        pe_cpu,
+                                    )
+                            # additional_information
+                            payload_info = getattr(nr, "additional_information", None)
+                            if payload_info is not None:
+                                info_dict = {}
+                                if isinstance(payload_info, dict):
+                                    # Already decoded
+                                    info_dict = payload_info
+                                elif isinstance(
+                                    payload_info, AdditionalInformationPayload
+                                ):
+                                    for k, entry in payload_info.entries.items():
+                                        if entry.tensor_data is not None:
+                                            dt = np.dtype(
+                                                getattr(
+                                                    entry, "tensor_dtype", "float32"
+                                                )
+                                            )
+                                            arr = np.frombuffer(
+                                                entry.tensor_data, dtype=dt
+                                            )
+                                            arr = arr.reshape(entry.tensor_shape)
+                                            info_dict[k] = torch.from_numpy(arr)
+                                        else:
+                                            info_dict[k] = entry.list_data
+                                if info_dict and req_id in self.requests:
+                                    setattr(
+                                        self.requests[req_id],
+                                        "additional_information_cpu",
+                                        info_dict,
+                                    )
+                except Exception as e:
+                    logger.error(
+                        f"Error decoding prompt_embeds / additional_information: {e}"
+                    )
+                    pass
+
+                if not scheduler_output.total_num_scheduled_tokens:
+                    if not has_kv_transfer_group():
+                        # Return empty ModelRunnerOutput if no work to do.
+                        return EMPTY_MODEL_RUNNER_OUTPUT
+                    return self.kv_connector_no_forward(
+                        scheduler_output, self.vllm_config
+                    )
+                if self.cache_config.kv_sharing_fast_prefill:
+                    assert not self.input_batch.num_prompt_logprobs, (
+                        "--kv-sharing-fast-prefill produces incorrect "
+                        "logprobs for prompt tokens, tokens, please disable "
+                        "it when the requests need prompt logprobs"
+                    )
+
+                # Prepare the decoder inputs.
+                (
+                    attn_metadata,
+                    logits_indices,
+                    spec_decode_metadata,
+                    num_scheduled_tokens_np,
+                    spec_decode_common_attn_metadata,
+                    max_query_len,
+                    ubatch_slices,
+                    num_tokens_after_padding,
+                ) = self._prepare_inputs(scheduler_output)
+
+            (
+                num_scheduled_tokens,
+                num_input_tokens,
+                num_tokens_across_dp,
+                input_ids,
+                inputs_embeds,
+                positions,
+                intermediate_tensors,
+                model_kwargs,
+                per_req_additional_information,
+            ) = self._preprocess(
+                scheduler_output,
+                num_scheduled_tokens_np,
+                intermediate_tensors,
+                ubatch_slices,
+                num_tokens_after_padding,
+            )
+
+            uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
+                num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
+            )
+            batch_descriptor = BatchDescriptor(
+                num_tokens=num_input_tokens, uniform_decode=uniform_decode
+            )
+            cudagraph_runtime_mode, batch_descriptor = (
+                self.cudagraph_dispatcher.dispatch(batch_descriptor)
+            )
+
+        # This is currently to get around the assert in the DPMetadata
+        # where it wants `num_tokens_across_dp` to align with `num_tokens`
+        if ubatch_slices is not None:
+            num_input_tokens = ubatch_slices[0].num_tokens
+
+        # Run the model.
+        # Use persistent buffers for CUDA graphs.
+        with (
+            set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
+                batch_descriptor=batch_descriptor,
+                ubatch_slices=ubatch_slices,
+            ),
+            record_function_or_nullcontext("Forward"),
+            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+        ):
+
+            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
+            # Always pass per-request runtime additional_information (persisted in request state)
+            try:
+                per_req_runtime_info = []
+                for req_id in self.input_batch.req_ids:
+                    req_state = self.requests.get(req_id)
+                    info = (
+                        getattr(req_state, "additional_information_cpu", None)
+                        if req_state is not None
+                        else None
+                    )
+                    per_req_runtime_info.append(info if isinstance(info, dict) else {})
+                model_kwargs_extra["runtime_additional_information"] = (
+                    per_req_runtime_info
+                )
+                model_kwargs_extra["request_ids"] = self.input_batch.req_ids
+                # Pass each request's token span within the flattened sequence for this step, enabling the model to map decode/prefill by request
+                req_token_spans = []
+                for req_index in range(len(self.input_batch.req_ids)):
+                    start_offset = int(self.query_start_loc.cpu[req_index])
+                    sched_tokens = int(num_scheduled_tokens_np[req_index])
+                    req_token_spans.append((start_offset, start_offset + sched_tokens))
+                model_kwargs_extra["request_token_spans"] = req_token_spans
+            except Exception:
+                pass
+            model_output = self.model(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+                **model_kwargs,
+                sampling_metadata=self.input_batch.sampling_metadata,
+                logits_index=logits_indices,
+                sampler=self.sampler,
+                **model_kwargs_extra,
+            )
+
+        with record_function_or_nullcontext("Postprocess"):
+            if self.use_aux_hidden_state_outputs:
+                # True when EAGLE 3 is used.
+                hidden_states, aux_hidden_states = model_output
+            else:
+                # Common case.
+                hidden_states = model_output
+                aux_hidden_states = None
+
+            hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+                hidden_states
+            )
+            # The model side may return per-request additional_information updates (model-agnostic channel).
+            # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
+            # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
+            try:
+                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(
+                        "additional_information_update"
+                    )
+                    if isinstance(updates_list, list):
+                        for idx, upd in enumerate(updates_list):
+                            if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
+                                continue
+                            req_id = self.input_batch.req_ids[idx]
+                            self._merge_additional_information_update(req_id, upd)
+                    # Option B: dict[str, dict] keyed by req_id
+                    updates_map = multimodal_outputs.get(
+                        "additional_information_update_by_req_id"
+                    )
+                    if isinstance(updates_map, dict):
+                        for req_id, upd in updates_map.items():
+                            if not isinstance(upd, dict):
+                                continue
+                            if req_id not in self.requests:
+                                continue
+                            self._merge_additional_information_update(req_id, upd)
+            except Exception as e:
+                logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}, with the multimodal_outputs as {multimodal_outputs}")
+            if not self.broadcast_pp_output:
+
+                # Common case.
+                if not get_pp_group().is_last_rank:
+                    # Return the intermediate tensors.
+                    assert isinstance(hidden_states, IntermediateTensors)
+                    hidden_states.kv_connector_output = kv_connector_output
+                    return hidden_states
+
+                if self.is_pooling_model:
+                    # Return the pooling output.
+                    output = self._pool(
+                        hidden_states, num_scheduled_tokens, num_scheduled_tokens_np
+                    )
+                    output.kv_connector_output = kv_connector_output
+                    return output
+
+                sample_hidden_states = hidden_states[logits_indices]
+                logits = self.model.compute_logits(sample_hidden_states)
+            else:
+                # Rare case.
+                assert not self.is_pooling_model
+
+                if not get_pp_group().is_last_rank:
+                    all_gather_tensors = {
+                        "residual": not is_residual_scattered_for_sp(
+                            self.vllm_config, num_input_tokens
+                        )
+                    }
+                    get_pp_group().send_tensor_dict(
+                        hidden_states.tensors,
+                        all_gather_group=get_tp_group(),
+                        all_gather_tensors=all_gather_tensors,
+                    )
+                    logits = None
+                else:
+                    sample_hidden_states = hidden_states[logits_indices]
+                    logits = self.model.compute_logits(sample_hidden_states)
+
+                model_output_broadcast_data = {}
+                if logits is not None:
+                    model_output_broadcast_data["logits"] = logits.contiguous()
+
+                model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
+                    model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
+                )
+                assert model_output_broadcast_data is not None
+                logits = model_output_broadcast_data["logits"]
+
+            # Apply structured output bitmasks if present
+            if scheduler_output.grammar_bitmask is not None:
+                apply_grammar_bitmask(
+                    scheduler_output, self.input_batch, logits, self.device
+                )
+
+        with record_function_or_nullcontext("Sample"):
+            sampler_output = self._sample(logits, spec_decode_metadata)
+
+        def propose_draft_token_ids(sampled_token_ids):
+            assert spec_decode_common_attn_metadata is not None
+            with record_function_or_nullcontext("Draft"):
+                self._draft_token_ids = self.propose_draft_token_ids(
+                    scheduler_output,
+                    sampled_token_ids,
+                    self.input_batch.sampling_metadata,
+                    hidden_states,
+                    sample_hidden_states,
+                    aux_hidden_states,
+                    spec_decode_metadata,
+                    spec_decode_common_attn_metadata,
+                )
+
+        use_padded_batch_for_eagle = (
+            self.speculative_config
+            and self.speculative_config.use_eagle()
+            and not self.speculative_config.disable_padded_drafter_batch
+        )
+        effective_drafter_max_model_len = self.max_model_len
+        if effective_drafter_max_model_len is None:
+            effective_drafter_max_model_len = self.model_config.max_model_len
+        if (
+            self.speculative_config
+            and self.speculative_config.draft_model_config is not None
+            and self.speculative_config.draft_model_config.max_model_len is not None
+        ):
+            effective_drafter_max_model_len = (
+                self.speculative_config.draft_model_config.max_model_len
+            )
+        input_fits_in_drafter = spec_decode_common_attn_metadata and (
+            spec_decode_common_attn_metadata.seq_lens.max()
+            + self.speculative_config.num_speculative_tokens
+            <= effective_drafter_max_model_len
+        )
+        if use_padded_batch_for_eagle and input_fits_in_drafter:
+            # EAGLE speculative decoding can use the GPU sampled tokens
+            # as inputs, and does not need to wait for bookkeeping to finish.
+            propose_draft_token_ids(sampler_output.sampled_token_ids)
+
+        with record_function_or_nullcontext("Bookkeep"):
+            (
+                num_nans_in_logits,
+                logprobs_lists,
+                valid_sampled_token_ids,
+                prompt_logprobs_dict,
+                req_ids_output_copy,
+                req_id_to_index_output_copy,
+                invalid_req_indices,
+            ) = self._bookkeeping_sync(
+                scheduler_output,
+                sampler_output,
+                logits,
+                hidden_states,
+                num_scheduled_tokens,
+            )
+
+        if (
+            self.speculative_config
+            and not use_padded_batch_for_eagle
+            and input_fits_in_drafter
+        ):
+            # ngram and other speculative decoding methods use the sampled
+            # tokens on the CPU, so they are run after bookkeeping.
+            propose_draft_token_ids(valid_sampled_token_ids)
+
+        with record_function_or_nullcontext("EPLB"):
+            self.eplb_step()
+
+        # Convert to per-request tensors on CPU
+        hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
+        pooler_output: list[Optional[torch.Tensor]] = []
+        prev_logits_index = 0
+        for logits_index in logits_indices:
+            pooler_output.append(
+                hidden_states_cpu[prev_logits_index : logits_index + 1]
+            )
+            prev_logits_index = logits_index + 1
+
+        output = OmniModelRunnerOutput(
+            req_ids=req_ids_output_copy,
+            req_id_to_index=req_id_to_index_output_copy,
+            sampled_token_ids=valid_sampled_token_ids,
+            logprobs=logprobs_lists,
+            prompt_logprobs_dict=prompt_logprobs_dict,
+            pooler_output=(
+                pooler_output
+                if self.vllm_config.model_config.engine_output_type != "text"
+                else None
+            ),
+            kv_connector_output=kv_connector_output,
+            num_nans_in_logits=num_nans_in_logits,
+        )
+
+        if not self.use_async_scheduling:
+            return output
+
+        return AsyncGPUModelRunnerOutput(
+            model_runner_output=output,
+            sampled_token_ids=sampler_output.sampled_token_ids,
+            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:
+            return
+        existing = getattr(req_state, "additional_information_cpu", {})
+        if not isinstance(existing, dict):
+            existing = {}
+        merged = dict(existing)
+        for k, v in upd.items():
+            if isinstance(v, torch.Tensor):
+                merged[k] = v.detach().to("cpu").contiguous()
+            elif isinstance(v, list):
+                new_list = []
+                for item in v:
+                    if isinstance(item, torch.Tensor):
+                        new_list.append(item.detach().to("cpu").contiguous())
+                    else:
+                        new_list.append(item)
+                merged[k] = new_list
+            else:
+                merged[k] = v
+        setattr(req_state, "additional_information_cpu", merged)

vllm_omni/worker/gpu/gpu_ar_worker.py

@@ -1,43 +1,57 @@
-from vllm.v1.worker.gpu_worker import Worker as GPUWorker
-import torch
-import os
 import gc
-from vllm.utils import GiB_bytes, MemorySnapshot
-from vllm.platforms import current_platform
-from vllm.worker.worker import _check_if_gpu_supports_dtype
-from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-from vllm.model_executor import set_random_seed
-from vllm.v1.utils import report_usage_stats
+import os
 
-from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+import torch
+
+from vllm.model_executor import set_random_seed
+from vllm.platforms import current_platform
+from vllm.utils import GiB_bytes, MemorySnapshot
+from vllm.v1.utils import report_usage_stats
+from vllm.v1.worker.gpu_worker import Worker as GPUWorker
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+from vllm_omni.worker.gpu.gpu_ar_model_runner import GPUARModelRunner
 
 
 class GPUARWorker(GPUWorker):
     def init_device(self):
         if self.device_config.device.type == "cuda":
-            # torch.distributed.all_reduce does not free the input tensor until
-            # the synchronization point. This causes the memory usage to grow
-            # as the number of all_reduce calls increases. This env var disables
-            # this behavior.
-            # Related issue:
-            # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
-            os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
-
             # This env var set by Ray causes exceptions with graph building.
             os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
             self.device = torch.device(f"cuda:{self.local_rank}")
             current_platform.set_device(self.device)
 
-            _check_if_gpu_supports_dtype(self.model_config.dtype)
+            current_platform.check_if_supports_dtype(self.model_config.dtype)
+
+            # Initialize the distributed environment BEFORE taking
+            # memory snapshot
+            # This ensures NCCL buffers are allocated before we measure
+            # available memory
+            init_worker_distributed_environment(
+                self.vllm_config,
+                self.rank,
+                self.distributed_init_method,
+                self.local_rank,
+                current_platform.dist_backend,
+            )
+
+            # Set random seed.
+            set_random_seed(self.model_config.seed)
+
+            # Now take memory snapshot after NCCL is initialized
             gc.collect()
             torch.cuda.empty_cache()
 
             # take current memory snapshot
             self.init_snapshot = MemorySnapshot()
-            self.requested_memory = (self.init_snapshot.total_memory *
-                                     self.cache_config.gpu_memory_utilization)
+            self.requested_memory = (
+                self.init_snapshot.total_memory
+                * self.cache_config.gpu_memory_utilization
+            )
             if self.init_snapshot.free_memory < self.requested_memory:
-                GiB = lambda b: round(b / GiB_bytes, 2)
+
+                def GiB(b):
+                    return round(b / GiB_bytes, 2)
+
                 raise ValueError(
                     f"Free memory on device "
                     f"({GiB(self.init_snapshot.free_memory)}/"
@@ -48,20 +62,13 @@ class GPUARWorker(GPUWorker):
                     f"utilization or reduce GPU memory used by other processes."
                 )
         else:
-            raise RuntimeError(
-                f"Not support device type: {self.device_config.device}")
-        # Initialize the distributed environment.
-        init_worker_distributed_environment(self.vllm_config, self.rank,
-                                            self.distributed_init_method,
-                                            self.local_rank,
-                                            current_platform.dist_backend)
-        # Set random seed.
-        set_random_seed(self.model_config.seed)
+            raise RuntimeError(f"Not support device type: {self.device_config.device}")
 
         # Construct the model runner
         self.model_runner: GPUARModelRunner = GPUARModelRunner(
-            self.vllm_config, self.device)
+            self.vllm_config, self.device
+        )
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.
-            report_usage_stats(self.vllm_config)
+            report_usage_stats(self.vllm_config)

vllm_omni/worker/gpu/gpu_diffusion_model_runner.py

@@ -0,0 +1,620 @@
+from __future__ import annotations
+
+import gc
+import logging
+from typing import Any, List, Optional, Union
+
+import numpy as np
+import torch
+
+from vllm.config import CUDAGraphMode
+from vllm.forward_context import BatchDescriptor
+from vllm.multimodal.inputs import MultiModalKwargs
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.spec_decode.eagle import EagleProposer
+from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker.gpu_model_runner import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    AsyncGPUModelRunnerOutput,
+    IntermediateTensors,
+    PerLayerAttnMetadata,
+    get_pp_group,
+    set_forward_context,
+)
+from vllm.v1.worker.ubatch_utils import UBatchSlices
+from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.gpu.gpu_model_runner import OmniGPUModelRunner
+
+logger = logging.getLogger(__name__)
+
+
+class GPUDiffusionModelRunner(OmniGPUModelRunner):
+    """Diffusion model runner for vLLM-omni (non-autoregressive).
+
+    - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
+    - Does not compute logits or perform token sampling.
+    - Executes diffusion process and returns tensors via `pooler_output`.
+    """
+
+    def _preprocess(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        ubatch_slices: Optional[UBatchSlices] = None,
+        num_tokens_after_padding: Optional[torch.Tensor] = None,
+    ) -> tuple[
+        int,
+        int,
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        torch.Tensor,
+        Optional[IntermediateTensors],
+        dict[str, Any],
+    ]:
+
+        # Input token count for this iteration (not used by diffusion, but
+        # retained to keep DP padding/ordering consistent)
+        num_input_tokens = scheduler_output.total_num_scheduled_tokens
+        num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
+        num_input_tokens += num_pad
+
+        # _prepare_inputs may reorder the batch, so we must gather multi
+        # modal outputs after that to ensure the correct order
+        if (
+            self.supports_mm_inputs
+            and get_pp_group().is_first_rank
+            and not self.model_config.is_encoder_decoder
+        ):
+            # Run the multimodal encoder if any.
+            self._execute_mm_encoder(scheduler_output)
+            mm_embeds = self._gather_mm_embeddings(scheduler_output)
+
+            # NOTE(woosuk): To unify token ids and soft tokens (vision
+            # embeddings), we always use embeddings (rather than token ids)
+            # as input to the multimodal model, even when the input is text.
+            inputs_embeds_scheduled = self.model.get_input_embeddings(
+                input_ids=self.input_ids.gpu[:num_input_tokens],
+                multimodal_embeddings=mm_embeds or None,
+            )
+
+            # TODO(woosuk): Avoid the copy. Optimize.
+            self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled)
+
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = {
+                **self._init_model_kwargs(num_input_tokens),
+                **self._extract_mm_kwargs(scheduler_output),
+            }
+        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
+            # Get the input embeddings for the tokens that are not input embeds,
+            # then put them into the appropriate positions.
+            # TODO(qthequartermasterman): Since even when prompt embeds are
+            # enabled, (a) not all requests will use prompt embeds, and (b)
+            # after the initial prompt is processed, the rest of the generated
+            # tokens will be token ids, it is not desirable to have the
+            # embedding layer outside of the CUDA graph all the time. The v0
+            # engine avoids this by "double compiling" the CUDA graph, once
+            # with input_ids and again with inputs_embeds, for all num_tokens.
+            # If a batch only has token ids, then including the embedding layer
+            # in the CUDA graph will be more performant (like in the else case
+            # below).
+            token_ids_idx = (
+                self.is_token_ids.gpu[:num_input_tokens]
+                .nonzero(as_tuple=False)
+                .squeeze(1)
+            )
+            # Some tokens ids may need to become embeds
+            if token_ids_idx.numel() > 0:
+                token_ids = self.input_ids.gpu[token_ids_idx]
+                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
+                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
+
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+        else:
+            # For text-only models, we use token ids as input.
+            # While it is possible to use embeddings as input just like the
+            # multimodal models, it is not desirable for performance since
+            # then the embedding layer is not included in the CUDA graph.
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = None
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+        if self.uses_mrope:
+            positions = self.mrope_positions.gpu[:, :num_input_tokens]
+        else:
+            positions = self.positions.gpu[:num_input_tokens]
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True
+            )
+
+        if (
+            self.model_config.is_encoder_decoder
+            and scheduler_output.scheduled_encoder_inputs
+        ):
+            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
+            model_kwargs.update(encoder_inputs)
+
+        return (
+            num_input_tokens,
+            num_input_tokens,
+            num_tokens_after_padding,
+            input_ids,
+            inputs_embeds,
+            positions,
+            intermediate_tensors,
+            model_kwargs,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
+        with record_function_or_nullcontext("Preprocess"):
+            with self.synchronize_input_prep():
+                # Update persistent batch states.
+                self._update_states(scheduler_output)
+                if not scheduler_output.total_num_scheduled_tokens:
+                    return EMPTY_MODEL_RUNNER_OUTPUT
+
+                # Prepare the decoder inputs.
+                (
+                    attn_metadata,
+                    logits_indices,
+                    spec_decode_metadata,
+                    num_scheduled_tokens_np,
+                    spec_decode_common_attn_metadata,
+                    max_query_len,
+                    ubatch_slices,
+                    num_tokens_after_padding,
+                ) = self._prepare_inputs(scheduler_output)
+
+            # Input token count for this iteration (not used by diffusion, but
+            # retained to keep DP padding/ordering consistent)
+            (
+                num_scheduled_tokens,
+                num_input_tokens,
+                num_tokens_across_dp,
+                input_ids,
+                inputs_embeds,
+                positions,
+                intermediate_tensors,
+                model_kwargs,
+            ) = self._preprocess(
+                scheduler_output,
+                intermediate_tensors,
+                ubatch_slices,
+                num_tokens_after_padding,
+            )
+
+        cudagraph_runtime_mode = CUDAGraphMode.NONE
+        # Run the model.
+        # Use persistent buffers for CUDA graphs.
+        with (
+            set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
+                batch_descriptor=None,
+                ubatch_slices=ubatch_slices,
+            ),
+            record_function_or_nullcontext("Forward"),
+            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+        ):
+
+            outputs = self._run_diffusion(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+                multimodal_kwargs=model_kwargs,
+                logits_indices=logits_indices,
+            )
+
+        _, multimodal_outputs = self.extract_multimodal_outputs(outputs)
+        # Ensure one tensor per request, map to CPU for output struct
+        pooler_output: List[Optional[torch.Tensor]] = []
+        if isinstance(multimodal_outputs, torch.Tensor):
+            # If model returned a single stacked tensor, split by requests
+            assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
+            for i in range(self.input_batch.num_reqs):
+                pooler_output.append(
+                    multimodal_outputs[i].detach().to("cpu").contiguous()
+                )
+        elif isinstance(multimodal_outputs, list):
+            for out in multimodal_outputs:
+                pooler_output.append(
+                    out.detach().to("cpu").contiguous() if out is not None else None
+                )
+        elif isinstance(multimodal_outputs, dict):
+            for out in multimodal_outputs.values():
+                pooler_output.append(
+                    out.detach().to("cpu").contiguous() if out is not None else None
+                )
+        else:
+            raise RuntimeError("Unsupported diffusion output type")
+
+        output = OmniModelRunnerOutput(
+            req_ids=self.input_batch.req_ids,
+            req_id_to_index=self.input_batch.req_id_to_index,
+            sampled_token_ids=[],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=pooler_output,
+            kv_connector_output=kv_connector_output,
+            num_nans_in_logits={},
+        )
+
+        if not self.use_async_scheduling:
+            return output
+
+        return AsyncGPUModelRunnerOutput(
+            model_runner_output=output,
+            sampled_token_ids=[],
+            invalid_req_indices=[],
+            async_output_copy_stream=self.async_output_copy_stream,
+        )
+
+    def _run_diffusion(
+        self,
+        *,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors],
+        inputs_embeds: Optional[torch.Tensor],
+        multimodal_kwargs: dict,
+        logits_indices: torch.Tensor,
+    ) -> Union[torch.Tensor, list[torch.Tensor]]:
+        """Runs the diffusion process and returns per-request tensors.
+
+        Tries model interfaces in the following order for maximal compatibility:
+        1) model.sample(condition=..., **kwargs)
+        2) model.forward(condition=..., **kwargs)
+        3) model.diffuse(condition=..., **kwargs)
+        """
+        # Keep inputs identical to AR runner
+        kwargs = dict(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
+            sampling_metadata=self.input_batch.sampling_metadata,
+            logits_index=logits_indices,
+            sampler=self.sampler,
+        )
+
+        # For Qwen 2.5 Omni's current implementation, we only support the forward method
+        if hasattr(self.model, "forward"):
+            return self.model.forward(**kwargs)
+        # TODO: add the diffuse method for other models
+
+        raise RuntimeError(
+            "The loaded model does not expose diffusion interfaces 'sample', "
+            "'forward', or 'diffuse'. Please implement one of them or adapt the runner."
+        )
+
+    @torch.inference_mode()
+    def _dummy_run(
+        self,
+        num_tokens: int,
+        cudagraph_runtime_mode: Optional[CUDAGraphMode] = None,
+        force_attention: bool = False,
+        uniform_decode: bool = False,
+        skip_eplb: bool = False,
+        is_profile: bool = False,
+        create_mixed_batch: bool = False,
+        remove_lora: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Run a dummy forward pass to warm up/profile run or capture the
+        CUDA graph for the model.
+
+        Args:
+            num_tokens: Number of tokens to run the dummy forward pass.
+            cudagraph_runtime_mode: used to control the behavior.
+                - if not set will determine the cudagraph mode based on using
+                    the self.cudagraph_dispatcher.
+                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
+                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
+                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
+                    needed.
+            force_attention: If True, always create attention metadata. Used to
+                warm up attention backend when mode is NONE.
+            uniform_decode: If True, the batch is a uniform decode batch.
+            skip_eplb: If True, skip EPLB state update.
+            is_profile: If True, this is a profile run.
+            create_mixed_batch: If True, create a mixed batch with both decode
+                (1 token) and prefill (multiple tokens) requests.
+            remove_lora: If False, dummy LoRAs are not destroyed after the run
+        """
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
+            CUDAGraphMode.NONE,
+            CUDAGraphMode.PIECEWISE,
+            CUDAGraphMode.FULL,
+        }
+
+        # If cudagraph_mode.decode_mode() == FULL and
+        # cudagraph_mode.separate_routine(). This means that we are using
+        # different graphs and/or modes for mixed prefill-decode batches vs.
+        # uniform decode batches. A uniform decode batch means that all
+        # requests have identical query length, except a potential virtual
+        # request (shorter) in the batch account for padding.
+        # Uniform decode batch could either be common pure decode, where
+        # max_query_len == 1, or speculative decode, where
+        # max_query_len == 1 + num_spec_decode_tokens.
+
+        # When setting max_query_len = 1, we switch to and capture the optimized
+        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
+        # for GQA/MQA.
+        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
+
+        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
+        # for dummy run with LoRA so that the num_reqs collectively
+        # has num_tokens in total.
+        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+        max_num_reqs = self.scheduler_config.max_num_seqs
+
+        num_reqs = min(num_tokens, max_num_reqs)
+        min_tokens_per_req = num_tokens // num_reqs
+        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
+        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+
+        assert sum(num_scheduled_tokens_list) == num_tokens
+        assert len(num_scheduled_tokens_list) == num_reqs
+        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
+
+        num_tokens_after_padding = None
+
+        # If we failed to microbatch, currently need to resynchronize
+        # TODO(lucas,sage): we should be able to avoid this second sync by
+        #  refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
+        #  a single `coordinate_batch_across_dp` function.
+        if num_tokens_after_padding is None:
+            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
+            num_tokens_after_padding = num_tokens + num_pad
+        else:
+            num_tokens_across_dp = num_tokens_after_padding
+            num_tokens_after_padding = int(num_tokens_after_padding[0].item())
+
+        attn_metadata: Optional[PerLayerAttnMetadata] = None
+
+        # If force_attention is True, we always capture attention. Otherwise,
+        # it only happens for cudagraph_runtime_mode=FULL.
+        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
+            attn_metadata = {}
+
+            seq_lens = max_query_len
+            self.seq_lens.np[:num_reqs] = seq_lens
+            self.seq_lens.np[num_reqs:] = 0
+            self.seq_lens.copy_to_gpu()
+
+            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+            self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+            self.query_start_loc.copy_to_gpu()
+
+            for kv_cache_group_id, kv_cache_group_spec in enumerate(
+                self.kv_cache_config.kv_cache_groups
+            ):
+                common_attn_metadata = CommonAttentionMetadata(
+                    query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
+                    query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
+                    seq_lens=self.seq_lens.gpu[:num_reqs],
+                    seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
+                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
+                        :num_reqs
+                    ],
+                    num_reqs=num_reqs,
+                    num_actual_tokens=num_tokens,
+                    max_query_len=max_query_len,
+                    max_seq_len=self.max_model_len,
+                    block_table_tensor=self.input_batch.block_table[
+                        kv_cache_group_id
+                    ].get_device_tensor(num_reqs),
+                    slot_mapping=self.input_batch.block_table[
+                        kv_cache_group_id
+                    ].slot_mapping.gpu[:num_tokens],
+                    causal=True,
+                )
+                for attn_group in self.attn_groups[kv_cache_group_id]:
+
+                    assert type(attn_metadata) is dict
+                    attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(  # noqa: E501
+                        common_attn_metadata
+                    )
+                    for layer_name in attn_group.layer_names:
+                        attn_metadata[layer_name] = attn_metadata_i
+
+        with self.maybe_dummy_run_with_lora(
+            self.lora_config, num_scheduled_tokens, remove_lora
+        ):
+            model_kwargs = self._init_model_kwargs(num_tokens)
+            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                model_kwargs = {
+                    **model_kwargs,
+                    **self._dummy_mm_kwargs(num_reqs),
+                }
+            elif self.enable_prompt_embeds:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                model_kwargs = self._init_model_kwargs(num_tokens)
+            else:
+                input_ids = self.input_ids.gpu[:num_tokens]
+                inputs_embeds = None
+
+            if self.uses_mrope:
+                positions = self.mrope_positions.gpu[:, :num_tokens]
+            else:
+                positions = self.positions.gpu[:num_tokens]
+
+            if get_pp_group().is_first_rank:
+                intermediate_tensors = None
+            else:
+                if self.intermediate_tensors is None:
+                    self.intermediate_tensors = (
+                        self.model.make_empty_intermediate_tensors(
+                            batch_size=self.max_num_tokens,
+                            dtype=self.model_config.dtype,
+                            device=self.device,
+                        )
+                    )
+
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                    num_tokens, None, False
+                )
+
+            # filter out the valid batch descriptor
+            _cg_mode, batch_descriptor = (
+                self.cudagraph_dispatcher.dispatch(
+                    BatchDescriptor(
+                        num_tokens=num_tokens_after_padding,
+                        uniform_decode=uniform_decode,
+                    )
+                )
+                if not is_profile
+                else (CUDAGraphMode.NONE, None)
+            )
+            if cudagraph_runtime_mode is not None:
+                # we allow forcing NONE when the dispatcher disagrees to support
+                # warm ups for cudagraph capture
+                assert (
+                    cudagraph_runtime_mode == CUDAGraphMode.NONE
+                    or cudagraph_runtime_mode == _cg_mode
+                ), (
+                    f"Cudagraph runtime mode mismatch at dummy_run. "
+                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
+                )
+            else:
+                cudagraph_runtime_mode = _cg_mode
+
+            with self.maybe_randomize_inputs(input_ids), set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_tokens_after_padding,
+                num_tokens_across_dp=num_tokens_across_dp,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
+                batch_descriptor=batch_descriptor,
+            ):
+                outputs = self.model(
+                    input_ids=input_ids,
+                    positions=positions,
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                    **model_kwargs,
+                )
+
+            if self.use_aux_hidden_state_outputs:
+                hidden_states, _ = outputs
+            else:
+                hidden_states = outputs
+
+            if self.speculative_config and self.speculative_config.use_eagle():
+                assert isinstance(self.drafter, EagleProposer)
+                self.drafter.dummy_run(num_tokens)
+
+        # This is necessary to avoid blocking DP.
+        # For dummy runs, we typically skip EPLB since we don't have any real
+        # requests to process.
+        # However, in DP settings, there may be cases when some DP ranks do
+        # not have any requests to process, so they're executing dummy batches.
+        # In such cases, we still have to trigger EPLB to make sure
+        # ranks execute the rearrangement in synchronization.
+
+        # logit_indices = np.cumsum(num_scheduled_tokens) - 1
+        hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
+        return hidden_states, None
+
+    @torch.inference_mode()
+    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+        logger.warning("Dummy sampler run is not implemented for diffusion model")
+        return None
+
+    def profile_run(self) -> None:
+        # Profile with multimodal encoder & encoder cache.
+        if self.supports_mm_inputs:
+            if self.model_config.multimodal_config.skip_mm_profiling:
+                logger.info(
+                    "Skipping memory profiling for multimodal encoder and "
+                    "encoder cache."
+                )
+            else:
+                mm_budget = self.mm_budget
+                assert mm_budget is not None
+
+                if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
+                    # NOTE: Currently model is profiled with a single non-text
+                    # modality with the max possible input tokens even when
+                    # it supports multiple.
+                    dummy_modality = mm_budget.get_modality_with_max_tokens()
+                    max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[
+                        dummy_modality
+                    ]
+
+                    logger.info(
+                        "Encoder cache will be initialized with a budget of "
+                        "%s tokens, and profiled with %s %s items of the "
+                        "maximum feature size.",
+                        encoder_budget,
+                        max_mm_items_per_batch,
+                        dummy_modality,
+                    )
+
+                    # Create dummy batch of multimodal inputs.
+                    batched_dummy_mm_inputs = self._get_mm_dummy_batch(
+                        dummy_modality,
+                        max_mm_items_per_batch,
+                    )
+
+                    # Run multimodal encoder.
+                    dummy_encoder_outputs = self.model.get_multimodal_embeddings(
+                        **batched_dummy_mm_inputs
+                    )
+
+                    sanity_check_mm_encoder_outputs(
+                        dummy_encoder_outputs,
+                        expected_num_items=max_mm_items_per_batch,
+                    )
+
+                    # NOTE: This happens when encoder cache needs to store
+                    # the embeddings that encoder outputs are scattered onto.
+                    # In this case we create dummy embeddings of size
+                    # (encode_budget, hidden_size) and scatter encoder
+                    # output into it.
+                    encoder_output_shape = dummy_encoder_outputs[0].shape
+                    if encoder_output_shape[0] < encoder_budget:
+                        expanded_outputs = []
+                        for output in dummy_encoder_outputs:
+                            expanded = output.new_zeros(
+                                (encoder_budget, encoder_output_shape[-1])
+                            )
+                            num_tokens = output.shape[0]
+                            expanded[:num_tokens].copy_(output)
+                            expanded_outputs.append(expanded)
+
+                        dummy_encoder_outputs = expanded_outputs
+
+                    # Cache the dummy encoder outputs.
+                    self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
+
+        # Add `is_profile` here to pre-allocate communication buffers
+        hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
+        if get_pp_group().is_last_rank:
+            pass
+        self._sync_device()
+        del hidden_states
+        self.encoder_cache.clear()
+        gc.collect()

vllm_omni/worker/gpu/gpu_diffusion_worker.py

@@ -1,43 +1,57 @@
-from vllm.v1.worker.gpu_worker import Worker as GPUWorker
-import torch
-import os
 import gc
-from vllm.utils import GiB_bytes, MemorySnapshot
-from vllm.platforms import current_platform
-from vllm.worker.worker import _check_if_gpu_supports_dtype
-from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-from vllm.model_executor import set_random_seed
-from vllm.v1.utils import report_usage_stats
+import os
 
-from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
+import torch
+
+from vllm.model_executor import set_random_seed
+from vllm.platforms import current_platform
+from vllm.utils import GiB_bytes, MemorySnapshot
+from vllm.v1.utils import report_usage_stats
+from vllm.v1.worker.gpu_worker import Worker as GPUWorker
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+from vllm_omni.worker.gpu.gpu_diffusion_model_runner import GPUDiffusionModelRunner
 
 
 class GPUDiffusionWorker(GPUWorker):
     def init_device(self):
         if self.device_config.device.type == "cuda":
-            # torch.distributed.all_reduce does not free the input tensor until
-            # the synchronization point. This causes the memory usage to grow
-            # as the number of all_reduce calls increases. This env var disables
-            # this behavior.
-            # Related issue:
-            # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
-            os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
-
             # This env var set by Ray causes exceptions with graph building.
             os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
             self.device = torch.device(f"cuda:{self.local_rank}")
             current_platform.set_device(self.device)
 
-            _check_if_gpu_supports_dtype(self.model_config.dtype)
+            current_platform.check_if_supports_dtype(self.model_config.dtype)
+
+            # Initialize the distributed environment BEFORE taking
+            # memory snapshot
+            # This ensures NCCL buffers are allocated before we measure
+            # available memory
+            init_worker_distributed_environment(
+                self.vllm_config,
+                self.rank,
+                self.distributed_init_method,
+                self.local_rank,
+                current_platform.dist_backend,
+            )
+
+            # Set random seed.
+            set_random_seed(self.model_config.seed)
+
+            # Now take memory snapshot after NCCL is initialized
             gc.collect()
             torch.cuda.empty_cache()
 
             # take current memory snapshot
             self.init_snapshot = MemorySnapshot()
-            self.requested_memory = (self.init_snapshot.total_memory *
-                                     self.cache_config.gpu_memory_utilization)
+            self.requested_memory = (
+                self.init_snapshot.total_memory
+                * self.cache_config.gpu_memory_utilization
+            )
             if self.init_snapshot.free_memory < self.requested_memory:
-                GiB = lambda b: round(b / GiB_bytes, 2)
+
+                def GiB(b):
+                    return round(b / GiB_bytes, 2)
+
                 raise ValueError(
                     f"Free memory on device "
                     f"({GiB(self.init_snapshot.free_memory)}/"
@@ -48,20 +62,13 @@ class GPUDiffusionWorker(GPUWorker):
                     f"utilization or reduce GPU memory used by other processes."
                 )
         else:
-            raise RuntimeError(
-                f"Not support device type: {self.device_config.device}")
-        # Initialize the distributed environment.
-        init_worker_distributed_environment(self.vllm_config, self.rank,
-                                            self.distributed_init_method,
-                                            self.local_rank,
-                                            current_platform.dist_backend)
-        # Set random seed.
-        set_random_seed(self.model_config.seed)
+            raise RuntimeError(f"Not support device type: {self.device_config.device}")
 
         # Construct the model runner
         self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(
-            self.vllm_config, self.device)
+            self.vllm_config, self.device
+        )
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.
-            report_usage_stats(self.vllm_config)
+            report_usage_stats(self.vllm_config)

vllm_omni/worker/gpu/gpu_model_runner.py

@@ -0,0 +1,588 @@
+from typing import TYPE_CHECKING, List, Optional, Union, cast
+
+import numpy as np
+import torch
+
+from vllm.config import CUDAGraphMode
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.forward_context import BatchDescriptor, set_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces_base import VllmModelForPooling
+from vllm.sampling_params import SamplingType
+from vllm.utils import LazyLoader, cdiv
+from vllm.v1.attention.backends.utils import (
+    CommonAttentionMetadata,
+    split_attn_metadata,
+)
+from vllm.v1.spec_decode.eagle import EagleProposer
+from vllm.v1.worker.gpu_input_batch import CachedRequestState
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata
+from vllm.v1.worker.ubatch_splitting import ubatch_split
+
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.output import SchedulerOutput
+else:
+    xgr = LazyLoader("xgr", globals(), "xgrammar")
+    xgr_torch_compile = LazyLoader(
+        "xgr_torch_compile",
+        globals(),
+        "xgrammar.kernels.apply_token_bitmask_inplace_torch_compile",
+    )
+
+logger = init_logger(__name__)
+
+
+class OmniGPUModelRunner(GPUModelRunner):
+    def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
+        """Update the cached states and the persistent batch with the scheduler
+        output.
+
+        The updated states are used by the `_prepare_inputs` function to create
+        the input GPU tensors for the model.
+
+        The SamplingMetadata is updated and copied to the GPU if there is a
+        new/resumed/paused/finished request in the batch.
+        """
+        # Remove finished requests from the cached states.
+        for req_id in scheduler_output.finished_req_ids:
+            self.requests.pop(req_id, None)
+        # Remove the finished requests from the persistent batch.
+        # NOTE(woosuk): There could be an edge case where finished_req_ids and
+        # scheduled_req_ids overlap. This happens when a request is aborted and
+        # then resubmitted with the same ID. In this case, we treat them as two
+        # distinct requests - clearing the cached states for the first request
+        # and handling the second as a new request.
+        for req_id in scheduler_output.finished_req_ids:
+            self.input_batch.remove_request(req_id)
+
+        # Free the cached encoder outputs.
+        for mm_hash in scheduler_output.free_encoder_mm_hashes:
+            self.encoder_cache.pop(mm_hash, None)
+
+        # Remove the unscheduled requests from the persistent batch.
+        # NOTE(woosuk): The unscheduled requests are either preempted requests
+        # or running requests that are not scheduled in this step. We remove
+        # them from the persistent batch but keep their cached states since
+        # they will be scheduled again sometime in the future.
+        scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
+        cached_req_ids = self.input_batch.req_id_to_index.keys()
+        unscheduled_req_ids = cached_req_ids - scheduled_req_ids
+        # NOTE(woosuk): The persistent batch optimization assumes that
+        # consecutive batches contain mostly the same requests. If batches
+        # have low request overlap (e.g., alternating between two distinct
+        # sets of requests), this optimization becomes very inefficient.
+        for req_id in unscheduled_req_ids:
+            self.input_batch.remove_request(req_id)
+
+        reqs_to_add: list[CachedRequestState] = []
+        # Add new requests to the cached states.
+        for new_req_data in scheduler_output.scheduled_new_reqs:
+            req_id = new_req_data.req_id
+            sampling_params = new_req_data.sampling_params
+            pooling_params = new_req_data.pooling_params
+
+            if (
+                sampling_params
+                and sampling_params.sampling_type == SamplingType.RANDOM_SEED
+            ):
+                generator = torch.Generator(device=self.device)
+                generator.manual_seed(sampling_params.seed)
+            else:
+                generator = None
+
+            if self.is_pooling_model:
+                assert pooling_params is not None
+                task = pooling_params.task
+                assert task is not None, "You did not set `task` in the API"
+
+                model = cast(VllmModelForPooling, self.get_model())
+                to_update = model.pooler.get_pooling_updates(task)
+                to_update.apply(pooling_params)
+
+            req_state = CachedRequestState(
+                req_id=req_id,
+                prompt_token_ids=new_req_data.prompt_token_ids,
+                prompt_embeds=new_req_data.prompt_embeds,
+                mm_features=new_req_data.mm_features,
+                sampling_params=sampling_params,
+                pooling_params=pooling_params,
+                generator=generator,
+                block_ids=new_req_data.block_ids,
+                num_computed_tokens=new_req_data.num_computed_tokens,
+                output_token_ids=[],
+                lora_request=new_req_data.lora_request,
+            )
+            self.requests[req_id] = req_state
+
+            # If prompt embeddings are provided, decode and attach to inter_data
+            try:
+                if getattr(new_req_data, "prompt_embeds", None) is not None:
+                    payload = new_req_data.prompt_embeds
+                    dtype = getattr(np, payload.dtype)
+                    arr = np.frombuffer(payload.data, dtype=dtype)
+                    arr = arr.reshape(payload.shape)
+                    pe_cpu = torch.from_numpy(arr)
+                    # Store temporarily on CPU; later moved to device in builder
+                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
+                    # Also replace payload with Tensor for user visibility in
+                    # scheduler_output
+                    try:
+                        new_req_data.prompt_embeds = pe_cpu  # type: ignore[assignment]
+                    except Exception:
+                        pass
+            except Exception as e:
+                logger.error(f"Error decoding prompt embeds: {e}")
+            # Decode additional_information payloads (dictionary)
+            try:
+                if getattr(new_req_data, "additional_information", None) is not None:
+                    payload_info = new_req_data.additional_information
+                    info_dict = {}
+                    if isinstance(payload_info, dict):
+                        info_dict = payload_info
+                    else:
+                        from vllm_omni.engine import AdditionalInformationPayload
+
+                        if isinstance(payload_info, AdditionalInformationPayload):
+                            for k, entry in payload_info.entries.items():
+                                if entry.tensor_data is not None:
+                                    dt = np.dtype(
+                                        getattr(entry, "tensor_dtype", "float32")
+                                    )
+                                    arr = np.frombuffer(entry.tensor_data, dtype=dt)
+                                    arr = arr.reshape(entry.tensor_shape)
+                                    info_dict[k] = torch.from_numpy(arr)
+                                else:
+                                    info_dict[k] = entry.list_data
+                    if info_dict:
+                        setattr(
+                            self.requests[req_id],
+                            "additional_information_cpu",
+                            info_dict,
+                        )
+            except Exception as e:
+                logger.error(f"Error decoding additional information: {e}")
+                pass
+
+            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
+            if self.uses_mrope:
+                self._init_mrope_positions(req_state)
+
+            reqs_to_add.append(self.requests[req_id])
+
+        # Update the states of the running/resumed requests.
+        is_last_rank = get_pp_group().is_last_rank
+        req_data = scheduler_output.scheduled_cached_reqs
+        for i, req_id in enumerate(req_data.req_ids):
+            req_state = self.requests[req_id]
+            num_computed_tokens = req_data.num_computed_tokens[i]
+            new_block_ids = req_data.new_block_ids[i]
+            resumed_from_preemption = req_data.resumed_from_preemption[i]
+
+            # Update the cached states.
+            req_state.num_computed_tokens = num_computed_tokens
+
+            if not is_last_rank:
+                # When using PP, the scheduler sends the sampled tokens back,
+                # because there's no direct communication between the first-
+                # stage worker and the last-stage worker.
+                new_token_ids = req_data.new_token_ids[i]
+                # Add the sampled token(s) from the previous step (if any).
+                # This doesn't include "unverified" tokens like spec tokens.
+                num_new_tokens = (
+                    num_computed_tokens + len(new_token_ids) - req_state.num_tokens
+                )
+                if num_new_tokens == 1:
+                    # Avoid slicing list in most common case.
+                    req_state.output_token_ids.append(new_token_ids[-1])
+                elif num_new_tokens > 0:
+                    req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
+
+            # Update the block IDs.
+            if not resumed_from_preemption:
+                if new_block_ids is not None:
+                    # Append the new blocks to the existing block IDs.
+                    for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
+                        block_ids.extend(new_ids)
+            else:
+                assert new_block_ids is not None
+                # The request is resumed from preemption.
+                # Replace the existing block IDs with the new ones.
+                req_state.block_ids = new_block_ids
+
+            req_index = self.input_batch.req_id_to_index.get(req_id)
+            if req_index is None:
+                # The request is not in the persistent batch.
+                # The request was either preempted and resumed later, or was not
+                # scheduled in the previous step and needs to be added again.
+                reqs_to_add.append(req_state)
+                continue
+
+            # Update the persistent batch.
+            self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens
+            if new_block_ids is not None:
+                self.input_batch.block_table.append_row(new_block_ids, req_index)
+
+            # For the last rank, we don't need to update the token_ids_cpu
+            # because the sampled tokens are already cached.
+            if not is_last_rank:
+                # Add new_token_ids to token_ids_cpu.
+                start_token_index = num_computed_tokens
+                end_token_index = num_computed_tokens + len(new_token_ids)
+                self.input_batch.token_ids_cpu[
+                    req_index, start_token_index:end_token_index
+                ] = new_token_ids
+                self.input_batch.num_tokens_no_spec[req_index] = end_token_index
+                self.input_batch.num_tokens[req_index] = end_token_index
+
+            # Add spec_token_ids to token_ids_cpu.
+            spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
+                req_id, ()
+            )
+            if spec_token_ids:
+                num_spec_tokens = len(spec_token_ids)
+                start_index = self.input_batch.num_tokens_no_spec[req_index]
+                end_token_index = start_index + num_spec_tokens
+                self.input_batch.token_ids_cpu[
+                    req_index, start_index:end_token_index
+                ] = spec_token_ids
+                # NOTE(woosuk): `num_tokens` here may include spec tokens.
+                self.input_batch.num_tokens[req_index] += num_spec_tokens
+
+        # Add the new or resumed requests to the persistent batch.
+        # The smaller empty indices are filled first.
+        for request in reqs_to_add:
+            self.input_batch.add_request(request)
+
+        # Condense the batched states if there are gaps left by removed requests
+        self.input_batch.condense()
+        # Allow attention backend to reorder the batch, potentially
+        self._may_reorder_batch(scheduler_output)
+        # Refresh batch metadata with any pending updates.
+        self.input_batch.refresh_metadata()
+
+    @torch.inference_mode()
+    def extract_multimodal_outputs(
+        self, hidden_states: Union[torch.Tensor, List[torch.Tensor]]
+    ) -> dict:
+        if (
+            hasattr(self.model, "have_multimodal_outputs")
+            and self.model.have_multimodal_outputs
+        ):
+            text_hidden_states = hidden_states.text_hidden_states
+            multimodal_outputs = hidden_states.multimodal_outputs
+
+        elif isinstance(hidden_states, torch.Tensor):
+            text_hidden_states = hidden_states
+            multimodal_outputs = {}
+        elif isinstance(hidden_states, List):
+            text_hidden_states = hidden_states[0]
+            multimodal_outputs = {}
+        else:
+            raise ValueError(f"Invalid hidden states type: {type(hidden_states)}")
+        return text_hidden_states, multimodal_outputs
+
+    @torch.inference_mode()
+    def _dummy_run(
+        self,
+        num_tokens: int,
+        cudagraph_runtime_mode: Optional[CUDAGraphMode] = None,
+        force_attention: bool = False,
+        uniform_decode: bool = False,
+        allow_microbatching: bool = True,
+        skip_eplb: bool = False,
+        is_profile: bool = False,
+        create_mixed_batch: bool = False,
+        remove_lora: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Run a dummy forward pass to warm up/profile run or capture the
+        CUDA graph for the model.
+
+        Args:
+            num_tokens: Number of tokens to run the dummy forward pass.
+            cudagraph_runtime_mode: used to control the behavior.
+                - if not set will determine the cudagraph mode based on using
+                    the self.cudagraph_dispatcher.
+                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
+                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
+                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
+                    needed.
+            force_attention: If True, always create attention metadata. Used to
+                warm up attention backend when mode is NONE.
+            uniform_decode: If True, the batch is a uniform decode batch.
+            skip_eplb: If True, skip EPLB state update.
+            is_profile: If True, this is a profile run.
+            create_mixed_batch: If True, create a mixed batch with both decode
+                (1 token) and prefill (multiple tokens) requests.
+            remove_lora: If False, dummy LoRAs are not destroyed after the run
+        """
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
+            CUDAGraphMode.NONE,
+            CUDAGraphMode.PIECEWISE,
+            CUDAGraphMode.FULL,
+        }
+
+        # If cudagraph_mode.decode_mode() == FULL and
+        # cudagraph_mode.separate_routine(). This means that we are using
+        # different graphs and/or modes for mixed prefill-decode batches vs.
+        # uniform decode batches. A uniform decode batch means that all
+        # requests have identical query length, except a potential virtual
+        # request (shorter) in the batch account for padding.
+        # Uniform decode batch could either be common pure decode, where
+        # max_query_len == 1, or speculative decode, where
+        # max_query_len == 1 + num_spec_decode_tokens.
+
+        # When setting max_query_len = 1, we switch to and capture the optimized
+        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
+        # for GQA/MQA.
+        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
+
+        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
+        # for dummy run with LoRA so that the num_reqs collectively
+        # has num_tokens in total.
+        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+        max_num_reqs = self.scheduler_config.max_num_seqs
+        if create_mixed_batch:
+            assert not uniform_decode
+            # Create mixed batch:
+            # first half decode tokens, second half one prefill
+            num_decode_tokens = num_tokens // 2
+            num_prefill_tokens = num_tokens - num_decode_tokens
+            num_reqs = num_decode_tokens + 1
+
+            # Create decode requests (1 token each) followed by prefill request
+            num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens]
+            # Note: Overriding max_query_len to be the prefill tokens
+            max_query_len = num_prefill_tokens
+        elif uniform_decode:
+            assert not create_mixed_batch
+            num_reqs = cdiv(num_tokens, max_query_len)
+            num_scheduled_tokens_list = [max_query_len] * num_reqs
+            if num_tokens % max_query_len != 0:
+                num_scheduled_tokens_list[-1] = num_tokens % max_query_len
+        else:
+            num_reqs = min(num_tokens, max_num_reqs)
+            min_tokens_per_req = num_tokens // num_reqs
+            num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
+            num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+
+        assert sum(num_scheduled_tokens_list) == num_tokens
+        assert len(num_scheduled_tokens_list) == num_reqs
+        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
+
+        total_num_scheduled_tokens = int(num_scheduled_tokens.sum())
+
+        ubatch_slices = None
+        num_tokens_after_padding = None
+
+        # We currently only microbatch if the number of tokens is
+        # over a certain threshold.
+        if self.parallel_config.enable_dbo and allow_microbatching:
+            ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split(
+                num_scheduled_tokens,
+                total_num_scheduled_tokens,
+                total_num_scheduled_tokens,
+                uniform_decode=uniform_decode,
+                vllm_config=self.vllm_config,
+            )
+            # Currently when DBO is enabled `ubatch_split` returns
+            # the num_tokens_after_padding for a single ubatch, but we have 2
+            # TODO(sage,lucas): this is cruft that should be addressed in the
+            # padding refactor.
+            if ubatch_num_tokens_after_padding is not None:
+                num_tokens_after_padding = ubatch_num_tokens_after_padding * 2
+
+        # If we failed to microbatch, currently need to resynchronize
+        # TODO(lucas,sage): we should be able to avoid this second sync by
+        #  refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
+        #  a single `coordinate_batch_across_dp` function.
+        if num_tokens_after_padding is None:
+            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
+            num_tokens_after_padding = num_tokens + num_pad
+        else:
+            num_tokens_across_dp = num_tokens_after_padding
+            num_tokens_after_padding = int(num_tokens_after_padding[0].item())
+
+        attn_metadata: Optional[PerLayerAttnMetadata] = None
+
+        # If force_attention is True, we always capture attention. Otherwise,
+        # it only happens for cudagraph_runtime_mode=FULL.
+        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
+            attn_metadata = {}
+            if ubatch_slices is not None:
+                attn_metadata = [dict() for _ in range(len(ubatch_slices))]
+
+            if create_mixed_batch:
+                # In the mixed batch mode (used for FI warmup), we use
+                # shorter sequence lengths to run faster.
+                # TODO(luka) better system for describing dummy batches
+                seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
+            else:
+                seq_lens = max_query_len
+            self.seq_lens.np[:num_reqs] = seq_lens
+            self.seq_lens.np[num_reqs:] = 0
+            self.seq_lens.copy_to_gpu()
+
+            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+            self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+            self.query_start_loc.copy_to_gpu()
+
+            for kv_cache_group_id, kv_cache_group_spec in enumerate(
+                self.kv_cache_config.kv_cache_groups
+            ):
+                common_attn_metadata = CommonAttentionMetadata(
+                    query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
+                    query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
+                    seq_lens=self.seq_lens.gpu[:num_reqs],
+                    seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
+                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
+                        :num_reqs
+                    ],
+                    num_reqs=num_reqs,
+                    num_actual_tokens=num_tokens,
+                    max_query_len=max_query_len,
+                    max_seq_len=self.max_model_len,
+                    block_table_tensor=self.input_batch.block_table[
+                        kv_cache_group_id
+                    ].get_device_tensor(num_reqs),
+                    slot_mapping=self.input_batch.block_table[
+                        kv_cache_group_id
+                    ].slot_mapping.gpu[:num_tokens],
+                    causal=True,
+                )
+                for attn_group in self.attn_groups[kv_cache_group_id]:
+                    if ubatch_slices is not None:
+                        common_attn_metadata_list = split_attn_metadata(
+                            ubatch_slices, common_attn_metadata
+                        )
+                        for ubid, common_attn_metadata in enumerate(
+                            common_attn_metadata_list
+                        ):
+                            assert common_attn_metadata.max_query_len == 1
+                            attn_metadata_i = attn_group.get_metadata_builder(
+                                ubatch_id=ubid
+                            ).build_for_cudagraph_capture(common_attn_metadata)
+                            for layer_name in attn_group.layer_names:
+                                assert type(attn_metadata) is list
+                                attn_metadata[ubid][layer_name] = attn_metadata_i
+                    else:
+                        assert type(attn_metadata) is dict
+                        attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(  # noqa: E501
+                            common_attn_metadata
+                        )
+                        for layer_name in attn_group.layer_names:
+                            attn_metadata[layer_name] = attn_metadata_i
+
+        with self.maybe_dummy_run_with_lora(
+            self.lora_config, num_scheduled_tokens, remove_lora
+        ):
+            model_kwargs = self._init_model_kwargs(num_tokens)
+            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                model_kwargs = {
+                    **model_kwargs,
+                    **self._dummy_mm_kwargs(num_reqs),
+                }
+            elif self.enable_prompt_embeds:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                model_kwargs = self._init_model_kwargs(num_tokens)
+            else:
+                input_ids = self.input_ids.gpu[:num_tokens]
+                inputs_embeds = None
+
+            if self.uses_mrope:
+                positions = self.mrope_positions.gpu[:, :num_tokens]
+            else:
+                positions = self.positions.gpu[:num_tokens]
+
+            if get_pp_group().is_first_rank:
+                intermediate_tensors = None
+            else:
+                if self.intermediate_tensors is None:
+                    self.intermediate_tensors = (
+                        self.model.make_empty_intermediate_tensors(
+                            batch_size=self.max_num_tokens,
+                            dtype=self.model_config.dtype,
+                            device=self.device,
+                        )
+                    )
+
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                    num_tokens, None, False
+                )
+
+            # filter out the valid batch descriptor
+            _cg_mode, batch_descriptor = (
+                self.cudagraph_dispatcher.dispatch(
+                    BatchDescriptor(
+                        num_tokens=num_tokens_after_padding,
+                        uniform_decode=uniform_decode,
+                    )
+                )
+                if not is_profile
+                else (CUDAGraphMode.NONE, None)
+            )
+            if cudagraph_runtime_mode is not None:
+                # we allow forcing NONE when the dispatcher disagrees to support
+                # warm ups for cudagraph capture
+                assert (
+                    cudagraph_runtime_mode == CUDAGraphMode.NONE
+                    or cudagraph_runtime_mode == _cg_mode
+                ), (
+                    f"Cudagraph runtime mode mismatch at dummy_run. "
+                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
+                )
+            else:
+                cudagraph_runtime_mode = _cg_mode
+
+            if ubatch_slices is not None:
+                # Adjust values to reflect a single ubatch.
+                # TODO(sage,lucas): this is cruft that should be addressed in
+                #  the padding refactor.
+                num_tokens_after_padding = ubatch_slices[0].num_tokens
+                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,
+            ):
+                outputs = self.model(
+                    input_ids=input_ids,
+                    positions=positions,
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                    **model_kwargs,
+                )
+
+            if self.use_aux_hidden_state_outputs:
+                hidden_states, _ = outputs
+            else:
+                hidden_states = outputs
+
+            if self.speculative_config and self.speculative_config.use_eagle():
+                assert isinstance(self.drafter, EagleProposer)
+                self.drafter.dummy_run(num_tokens)
+
+        # This is necessary to avoid blocking DP.
+        # For dummy runs, we typically skip EPLB since we don't have any real
+        # requests to process.
+        # However, in DP settings, there may be cases when some DP ranks do
+        # not have any requests to process, so they're executing dummy batches.
+        # In such cases, we still have to trigger EPLB to make sure
+        # ranks execute the rearrangement in synchronization.
+        if not skip_eplb:
+            self.eplb_step(is_dummy=True, is_profile=is_profile)
+
+        logit_indices = np.cumsum(num_scheduled_tokens) - 1
+        hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+            hidden_states
+        )
+        return hidden_states, hidden_states[logit_indices]

vllm_omni/worker/gpu_ar_model_runner.py

@@ -1,578 +0,0 @@
-"""AR GPU Model Runner for vLLM-omni.
-
-Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
-and also outputs sampled tokens.
-"""
-
-from __future__ import annotations
-
-from typing import Optional, Union, Any, List
-import numpy as np
-
-import torch
-
-from vllm import envs
-from vllm.v1.worker.gpu_model_runner import (
-    EMPTY_MODEL_RUNNER_OUTPUT,
-    IntermediateTensors,
-    get_pp_group,
-    get_tp_group,
-    has_kv_transfer_group,
-    set_forward_context,
-)
-from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.multimodal.inputs import MultiModalKwargs
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
-from vllm.v1.spec_decode.eagle import EagleProposer
-
-from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload
-from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
-
-
-class GPUARModelRunner(OmniGPUModelRunner):
-    """Autoregressive GPU model runner that returns hidden states per request.
-
-    This runner follows the same preparation and forward path as GPUModelRunner
-    (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs),
-    and additionally performs lightweight sampling so that sampled tokens are
-    available in outputs. Hidden representations are taken at the same indices
-    that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`).
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.input_ids = torch.zeros(self.max_num_tokens,
-                                     dtype=torch.int64,
-                                     device=self.device)
-
-    @torch.inference_mode()
-    def execute_model(
-        self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-    ) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
-        # Update internal state with the new schedule
-        self._update_states(scheduler_output)
-
-        # Decode per-request prompt_embeds / additional_hidden_states payloads (if present) into CPU tensors
-        try:
-            new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
-            if new_reqs:
-                for nr in new_reqs:
-                    req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
-                    if req_id is None:
-                        continue
-                    # prompt_embeds
-                    payload_pe = getattr(nr, "prompt_embeds", None)
-                    if payload_pe is not None:
-                        if isinstance(payload_pe, torch.Tensor):
-                            pe_cpu = payload_pe.detach().to("cpu").contiguous()
-                        elif isinstance(payload_pe, PromptEmbedsPayload):
-                            dt = np.dtype(getattr(payload_pe, "dtype", "float32"))
-                            arr = np.frombuffer(payload_pe.data, dtype=dt)
-                            arr = arr.reshape(payload_pe.shape)
-                            pe_cpu = torch.from_numpy(arr)
-                        else:
-                            pe_cpu = None
-                        if pe_cpu is not None and req_id in self.requests:
-                            setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
-                    # additional_information
-                    payload_info = getattr(nr, "additional_information", None)
-                    if payload_info is not None:
-                        info_dict = {}
-                        if isinstance(payload_info, dict):
-                            # Already decoded
-                            info_dict = payload_info
-                        elif isinstance(payload_info, AdditionalInformationPayload):
-                            for k, entry in payload_info.entries.items():
-                                if entry.tensor_data is not None:
-                                    dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
-                                    arr = np.frombuffer(entry.tensor_data, dtype=dt)
-                                    arr = arr.reshape(entry.tensor_shape)
-                                    info_dict[k] = torch.from_numpy(arr)
-                                else:
-                                    info_dict[k] = entry.list_data
-                        if info_dict and req_id in self.requests:
-                            setattr(self.requests[req_id], "additional_information_cpu", info_dict)
-        except Exception:
-            pass
-
-        # If there's no work to do, either return empty output or kv-only path
-        if not scheduler_output.total_num_scheduled_tokens:
-            if not has_kv_transfer_group():
-                return EMPTY_MODEL_RUNNER_OUTPUT
-            return self.kv_connector_no_forward(scheduler_output,
-                                                self.vllm_config)
-
-        # Prepare decoder inputs and attention metadata
-        (attn_metadata, attention_cuda_graphs, logits_indices,
-         spec_decode_metadata, num_scheduled_tokens_np,
-         spec_decode_common_attn_metadata) = self._prepare_inputs(
-             scheduler_output)
-
-        # Determine number of input tokens for this iteration
-        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        if (self.use_cuda_graph
-                and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
-            num_input_tokens = self.vllm_config.pad_for_cudagraph(
-                num_scheduled_tokens)
-        else:
-            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
-            if (self.compilation_config.pass_config.enable_sequence_parallelism
-                    and tp_size > 1):
-                from vllm.utils import round_up  # lazy local import
-                num_input_tokens = round_up(num_scheduled_tokens, tp_size)
-            else:
-                num_input_tokens = num_scheduled_tokens
-
-        # DP padding
-        num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
-        num_input_tokens += num_pad
-
-        # Multimodal handling (encode and gather embeddings if needed)
-        if self.is_multimodal_model:
-            self._execute_mm_encoder(scheduler_output)
-            mm_embeds = self._gather_mm_embeddings(scheduler_output)
-        else:
-            mm_embeds = []
-
-        # Always assemble inputs_embeds on first PP rank; overlay per-request prompt_embeds and collect additional_hidden_states for prefill only
-        if get_pp_group().is_first_rank:
-            inputs_embeds_scheduled = self.model.get_input_embeddings(
-                input_ids=self.input_ids[:num_scheduled_tokens],
-                multimodal_embeddings=(mm_embeds or None)
-                if self.is_multimodal_model else None,
-            )
-
-            # Copy into persistent buffer to enable CUDA Graph capture
-            self.inputs_embeds[:num_scheduled_tokens].copy_(
-                inputs_embeds_scheduled)
-
-            # Reset per-step additional information collector
-            if hasattr(self, "_forward_additional_information"):
-                self._forward_additional_information = None
-
-            # Overlay custom prompt_embeds per request for the prompt portion; collect additional_information (tensor/list) for prefill portion only
-            for req_index, req_id in enumerate(self.input_batch.req_ids):
-                req_state = self.requests[req_id]
-                pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
-                addi_cpu = getattr(req_state, "additional_information_cpu", None)
-                num_computed_tokens = int(
-                    self.input_batch.num_computed_tokens_cpu[req_index])
-                prompt_len = len(req_state.prompt_token_ids)
-                prompt_remaining = max(0, prompt_len - num_computed_tokens)
-                sched_tokens = int(num_scheduled_tokens_np[req_index])
-                overlay_len = min(sched_tokens, prompt_remaining)
-                if overlay_len <= 0:
-                    continue
-                if pe_cpu is not None:
-                    src = pe_cpu[num_computed_tokens:
-                                 num_computed_tokens + overlay_len].to(
-                                      dtype=self.dtype,
-                                      device=self.device,
-                                      non_blocking=True)
-                    start_offset = int(self.query_start_loc_cpu[req_index])
-                    self.inputs_embeds[start_offset:start_offset + overlay_len] \
-                        .copy_(src)
-                # For additional_information: handle arbitrary keys
-                if addi_cpu is not None and isinstance(addi_cpu, dict):
-                    # Lazy init collector dict
-                    if not hasattr(self, "_forward_additional_information") or \
-                       self._forward_additional_information is None:
-                        self._forward_additional_information = {}
-                    # Process tensors (slice by scheduled prompt range) and lists (append per-request)
-                    for k, v in addi_cpu.items():
-                        if isinstance(v, torch.Tensor):
-                            # Slice along token dimension for prefill part
-                            try:
-                                seg = v[num_computed_tokens:
-                                        num_computed_tokens + overlay_len].to(
-                                            dtype=self.dtype,
-                                            device=self.device,
-                                            non_blocking=True)
-                            except Exception:
-                                # Fallback: move whole tensor if slicing fails
-                                seg = v.to(dtype=self.dtype,
-                                           device=self.device,
-                                           non_blocking=True)
-                            prev_val = self._forward_additional_information.get(k)
-                            self._forward_additional_information[k] = (
-                                torch.cat([prev_val, seg], dim=0)
-                                if isinstance(prev_val, torch.Tensor) else seg.clone())
-                        elif isinstance(v, list):
-                            prev_val = self._forward_additional_information.get(k)
-                            if prev_val is None:
-                                self._forward_additional_information[k] = [v]
-                            elif isinstance(prev_val, list):
-                                self._forward_additional_information[k].append(v)
-                            else:
-                                # Mixed types: wrap existing into list
-                                self._forward_additional_information[k] = [prev_val, v]
-
-            input_ids = self.input_ids[:num_input_tokens]  # preserved for APIs
-            inputs_embeds = self.inputs_embeds[:num_input_tokens]
-            model_mm_kwargs = (self._extract_mm_kwargs(scheduler_output)
-                               if self.is_multimodal_model else {})
-        else:
-            input_ids = self.input_ids[:num_input_tokens]
-            inputs_embeds = None
-            model_mm_kwargs = {}
-
-        # Positions/mRoPE
-        if self.uses_mrope:
-            positions = self.mrope_positions[:, :num_input_tokens]
-        else:
-            positions = self.positions[:num_input_tokens]
-
-        # Handle pipeline-parallel intermediate tensors
-        if get_pp_group().is_first_rank:
-            intermediate_tensors = None
-        else:
-            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True)
-
-        # Some attention backends only support CUDA Graphs in pure decode.
-        skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
-
-        # Forward pass
-        with set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_input_tokens,
-                num_tokens_across_dp=num_tokens_across_dp,
-                skip_cuda_graphs=skip_cuda_graphs,
-        ), self.maybe_get_kv_connector_output(
-                scheduler_output) as kv_connector_output:
-
-            model_kwargs_extra = {}
-            # Only pass additional_information for the prefill part
-            if hasattr(self, "_forward_additional_information") and \
-               self._forward_additional_information is not None and \
-               isinstance(self._forward_additional_information, dict):
-                model_kwargs_extra["additional_information"] = self._forward_additional_information
-            model_output = self.model(
-                input_ids=input_ids,
-                positions=positions,
-                intermediate_tensors=intermediate_tensors,
-                inputs_embeds=inputs_embeds,
-                **MultiModalKwargs.as_kwargs(
-                    model_mm_kwargs,
-                    device=self.device,
-                ),
-                sampling_metadata=self.input_batch.sampling_metadata,
-                logits_index=logits_indices,
-                sampler=self.sampler,
-                **model_kwargs_extra,
-            )
-
-        if self.use_aux_hidden_state_outputs:
-            hidden_states, _aux_hidden_states = model_output
-        else:
-            hidden_states = model_output
-            _aux_hidden_states = None
-
-        text_hidden_states, multimodal_outputs = (
-            self.extract_multimodal_outputs(hidden_states))
-
-        # Mid PP stages return intermediate tensors unmodified
-        if not get_pp_group().is_last_rank:
-            assert isinstance(text_hidden_states, IntermediateTensors)
-            text_hidden_states.kv_connector_output = kv_connector_output
-            return text_hidden_states
-
-        # Broadcast PP output for external_launcher (torchrun)
-        broadcast_pp_output = \
-            self.parallel_config.distributed_executor_backend \
-            == "external_launcher" and len(get_pp_group().ranks) > 0
-        if not get_pp_group().is_last_rank:
-            assert isinstance(text_hidden_states, IntermediateTensors)
-            if not broadcast_pp_output:
-                text_hidden_states.kv_connector_output = kv_connector_output
-                return text_hidden_states
-            get_pp_group().send_tensor_dict(text_hidden_states.tensors,
-                                            all_gather_group=get_tp_group())
-            logits = None
-        else:
-            if self.input_batch.pooling_params:
-                return self._pool(text_hidden_states, num_scheduled_tokens,
-                                  num_scheduled_tokens_np, kv_connector_output)
-
-            sample_hidden_states = text_hidden_states[logits_indices]
-            logits = self.model.compute_logits(sample_hidden_states, None)
-        if broadcast_pp_output:
-            model_output_broadcast_data = {
-                "logits": logits.contiguous(),
-            } if logits is not None else {}
-            model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
-                model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)
-            assert model_output_broadcast_data is not None
-            logits = model_output_broadcast_data["logits"]
-
-        # Apply structured output bitmasks if present
-        if scheduler_output.grammar_bitmask is not None:
-            self.apply_grammar_bitmask(scheduler_output, logits)
-
-        # Sample the next token and get logprobs if needed (with spec decode)
-        sampling_metadata = self.input_batch.sampling_metadata
-        if spec_decode_metadata is None:
-            sampler_output = self.sampler(
-                logits=logits,
-                sampling_metadata=sampling_metadata,
-            )
-        else:
-            assert logits is not None
-            bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
-            sampler_output = self.sampler(
-                logits=bonus_logits,
-                sampling_metadata=sampling_metadata,
-            )
-            bonus_token_ids = sampler_output.sampled_token_ids
-            target_logits = logits[spec_decode_metadata.target_logits_indices]
-            output_token_ids = self.rejection_sampler(
-                spec_decode_metadata,
-                None,
-                target_logits,
-                bonus_token_ids,
-                sampling_metadata,
-            )
-            sampler_output.sampled_token_ids = output_token_ids
-
-        num_nans_in_logits = {}
-        if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
-            num_nans_in_logits = self._get_nans_in_logits(logits)
-
-        # Handle partial prefill: discard sampled tokens and rewind RNG
-        discard_sampled_tokens_req_indices = []
-        for i, req_id in enumerate(self.input_batch.req_ids):
-            req_state = self.requests[req_id]
-            seq_len = (req_state.num_computed_tokens +
-                       scheduler_output.num_scheduled_tokens[req_id])
-            if seq_len < req_state.num_tokens:
-                generator = self.input_batch.generators.get(i)
-                if generator is not None:
-                    generator.set_offset(generator.get_offset() - 4)
-                discard_sampled_tokens_req_indices.append(i)
-
-        # Move CPU sync parts
-        logprobs_tensors = sampler_output.logprobs_tensors
-        logprobs_lists = logprobs_tensors.tolists() \
-            if logprobs_tensors is not None else None
-
-        # Prompt logprobs if needed
-        prompt_logprobs_dict = self._get_prompt_logprobs_dict(
-            text_hidden_states[:num_scheduled_tokens],
-            scheduler_output,
-        )
-
-        # Parse valid sampled tokens
-        sampled_token_ids = sampler_output.sampled_token_ids
-        max_gen_len = sampled_token_ids.shape[-1]
-        if max_gen_len == 1:
-            valid_sampled_token_ids = sampled_token_ids.tolist()
-        else:
-            valid_sampled_token_ids = self.rejection_sampler.parse_output(
-                sampled_token_ids,
-                self.input_batch.vocab_size,
-            )
-        for i in discard_sampled_tokens_req_indices:
-            valid_sampled_token_ids[i].clear()
-
-        # Cache sampled tokens
-        for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
-            if not sampled_ids:
-                continue
-            start_idx = self.input_batch.num_tokens_no_spec[req_idx]
-            end_idx = start_idx + len(sampled_ids)
-            self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids
-            self.input_batch.num_tokens_no_spec[req_idx] = end_idx
-            self.input_batch.num_tokens[req_idx] = end_idx
-            req_id = self.input_batch.req_ids[req_idx]
-            req_state = self.requests[req_id]
-            req_state.output_token_ids.extend(sampled_ids)
-
-        # Speculative decoding draft tokens if configured
-        if not self.speculative_config:
-            spec_token_ids = None
-        else:
-            assert spec_decode_common_attn_metadata is not None
-            spec_token_ids = self.propose_draft_token_ids(
-                scheduler_output,
-                valid_sampled_token_ids,
-                sampling_metadata,
-                text_hidden_states,
-                sample_hidden_states,
-                _aux_hidden_states,
-                spec_decode_metadata,
-                spec_decode_common_attn_metadata,
-            )
-
-        # Convert to per-request tensors on CPU
-        text_hidden_states_cpu = text_hidden_states.detach().to("cpu").contiguous()
-        pooler_output: list[Optional[torch.Tensor]] = []
-        prev_logits_index = 0
-        for logits_index in logits_indices:
-            pooler_output.append(text_hidden_states_cpu[prev_logits_index:logits_index+1])
-            prev_logits_index = logits_index + 1
-
-
-        self.eplb_step()
-
-        return OmniModelRunnerOutput(
-            req_ids=self.input_batch.req_ids,
-            req_id_to_index=self.input_batch.req_id_to_index,
-            sampled_token_ids=valid_sampled_token_ids,
-            spec_token_ids=spec_token_ids,
-            logprobs=logprobs_lists,
-            prompt_logprobs_dict=prompt_logprobs_dict,
-            pooler_output=pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None,
-            kv_connector_output=kv_connector_output,
-            num_nans_in_logits=num_nans_in_logits,
-        )
-    
-    @torch.inference_mode()
-    def extract_multimodal_outputs(self, hidden_states: Union[torch.Tensor, List[torch.Tensor]]) -> dict:
-        if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs:
-            text_hidden_states = hidden_states.text_hidden_states
-            multimodal_outputs = hidden_states.multimodal_outputs
-
-        elif isinstance(hidden_states, torch.Tensor):
-            text_hidden_states = hidden_states
-            multimodal_outputs = {}
-        elif isinstance(hidden_states, List):
-            text_hidden_states = hidden_states[0]
-            multimodal_outputs = {}
-        else:
-            raise ValueError(f"Invalid hidden states type: {type(hidden_states)}")
-        return text_hidden_states, multimodal_outputs
-    
-    @torch.inference_mode()
-    def _dummy_run(
-        self,
-        num_tokens: int,
-        capture_attn_cudagraph: bool = False,
-        skip_eplb: bool = False,
-        is_profile: bool = False,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-
-        # Padding for DP
-        num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
-        num_tokens += num_pad
-
-        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
-        # for dummy run with LoRA so that the num_reqs collectively
-        # has num_tokens in total.
-        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
-        max_num_reqs = self.scheduler_config.max_num_seqs
-        num_reqs = min(num_tokens, max_num_reqs)
-        min_tokens_per_req = num_tokens // num_reqs
-        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
-        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
-        assert sum(num_scheduled_tokens_list) == num_tokens
-        assert len(num_scheduled_tokens_list) == num_reqs
-        num_scheduled_tokens = np.array(num_scheduled_tokens_list,
-                                        dtype=np.int32)
-
-        attn_metadata: Optional[dict[str, Any]] = None
-        if capture_attn_cudagraph:
-            attn_metadata = {}
-
-            # Make sure max_model_len is used at the graph capture time.
-            self.seq_lens_np[:num_reqs] = self.max_model_len
-            self.seq_lens_np[num_reqs:] = 0
-            self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
-                                           non_blocking=True)
-
-            for kv_cache_group_id, kv_cache_group_spec in enumerate(
-                    self.kv_cache_config.kv_cache_groups):
-                common_attn_metadata = CommonAttentionMetadata(
-                    query_start_loc=self.query_start_loc[:num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs +
-                                                                 1],
-                    seq_lens=self.seq_lens[:num_reqs],
-                    seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.
-                    num_computed_tokens_cpu_tensor[:num_reqs],
-                    num_reqs=num_reqs,
-                    num_actual_tokens=num_tokens,
-                    max_query_len=num_tokens,
-                    block_table_tensor=self.input_batch.block_table[
-                        kv_cache_group_id].get_device_tensor()[:num_reqs],
-                    slot_mapping=self.input_batch.
-                    block_table[kv_cache_group_id].slot_mapping[:num_tokens],
-                    causal=True)
-
-                for attn_group in self.attn_groups[kv_cache_group_id]:
-                    attn_metadata_i = attn_group.metadata_builder\
-                        .build_for_cudagraph_capture(common_attn_metadata)
-                    for layer_name in kv_cache_group_spec.layer_names:
-                        attn_metadata[layer_name] = attn_metadata_i
-
-        with self.maybe_dummy_run_with_lora(self.lora_config,
-                                            num_scheduled_tokens):
-            if self.is_multimodal_model:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds[:num_tokens]
-                model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)
-            else:
-                input_ids = self.input_ids[:num_tokens]
-                inputs_embeds = None
-                model_mm_kwargs = {}
-
-            if self.uses_mrope:
-                positions = self.mrope_positions[:, :num_tokens]
-            else:
-                positions = self.positions[:num_tokens]
-
-            if get_pp_group().is_first_rank:
-                intermediate_tensors = None
-            else:
-                if self.intermediate_tensors is None:
-                    self.intermediate_tensors = (
-                        self.model.make_empty_intermediate_tensors(
-                            batch_size=self.max_num_tokens,
-                            dtype=self.model_config.dtype,
-                            device=self.device))
-
-                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                    num_tokens, None, False)
-
-            with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                    attn_metadata,
-                    self.vllm_config,
-                    num_tokens=num_tokens,
-                    num_tokens_across_dp=num_tokens_across_dp):
-                outputs = self.model(
-                    input_ids=input_ids,
-                    positions=positions,
-                    intermediate_tensors=intermediate_tensors,
-                    inputs_embeds=inputs_embeds,
-                    **MultiModalKwargs.as_kwargs(
-                        model_mm_kwargs,
-                        device=self.device,
-                    ),
-                )
-
-            if self.use_aux_hidden_state_outputs:
-                hidden_states, _ = outputs
-            else:
-                hidden_states = outputs
-
-            if self.speculative_config and self.speculative_config.use_eagle():
-                assert isinstance(self.drafter, EagleProposer)
-                self.drafter.dummy_run(num_tokens)
-
-        # This is necessary to avoid blocking DP.
-        # For dummy runs, we typically skip EPLB since we don't have any real
-        # requests to process.
-        # However, in DP settings, there may be cases when some DP ranks do
-        # not have any requests to process, so they're executing dummy batches.
-        # In such cases, we still have to trigger EPLB to make sure
-        # ranks execute the rearrangement in synchronization.
-        if not skip_eplb:
-            self.eplb_step(is_dummy=True, is_profile=is_profile)
-
-        logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
-        return hidden_states, hidden_states[logit_indices]
-

vllm_omni/worker/gpu_diffusion_model_runner.py

@@ -1,362 +0,0 @@
-from __future__ import annotations
-
-from typing import Optional, Union, List
-import numpy as np
-import gc
-import logging
-
-import torch
-
-from vllm.v1.worker.gpu_model_runner import (
-    EMPTY_MODEL_RUNNER_OUTPUT,
-    IntermediateTensors,
-    get_pp_group,
-    set_forward_context,
-)
-from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
-from vllm.multimodal.inputs import MultiModalKwargs
-
-from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
-
-from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
-
-
-logger = logging.getLogger(__name__)
-
-
-class GPUDiffusionModelRunner(OmniGPUModelRunner):
-    """Diffusion model runner for vLLM-omni (non-autoregressive).
-
-    - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
-    - Does not compute logits or perform token sampling.
-    - Executes diffusion process and returns tensors via `pooler_output`.
-    """
-
-    @torch.inference_mode()
-    def execute_model(
-        self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-    ) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
-        self._update_states(scheduler_output)
-
-        if not scheduler_output.total_num_scheduled_tokens:
-
-            return EMPTY_MODEL_RUNNER_OUTPUT
-           
-        # Prepare decoder inputs and attention metadata (for batch/order mapping)
-        (attn_metadata, attention_cuda_graphs, logits_indices,
-         spec_decode_metadata, num_scheduled_tokens_np,
-         spec_decode_common_attn_metadata) = self._prepare_inputs(
-             scheduler_output)
-
-        # Input token count for this iteration (not used by diffusion, but
-        # retained to keep DP padding/ordering consistent)
-        num_input_tokens = scheduler_output.total_num_scheduled_tokens
-        num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
-        num_input_tokens += num_pad
-
-        # Multimodal conditioning (e.g., text/audio/video encoders)
-        if self.is_multimodal_model:
-            self._execute_mm_encoder(scheduler_output)
-            mm_embeds = self._gather_mm_embeddings(scheduler_output)
-        else:
-            mm_embeds = []
-
-        # Build inputs to mirror AR runner: input_ids/positions/embeds
-        if self.is_multimodal_model and get_pp_group().is_first_rank:
-            inputs_embeds_scheduled = self.model.get_input_embeddings(
-                input_ids=self.input_ids[:scheduler_output.total_num_scheduled_tokens],
-                multimodal_embeddings=mm_embeds or None,
-            )
-            self.inputs_embeds[:scheduler_output.total_num_scheduled_tokens].copy_(
-                inputs_embeds_scheduled)
-            input_ids = self.input_ids[:num_input_tokens]
-            inputs_embeds = self.inputs_embeds[:num_input_tokens]
-            model_mm_kwargs = self._extract_mm_kwargs(scheduler_output)
-        else:
-            input_ids = self.input_ids[:num_input_tokens]
-            inputs_embeds = None
-            model_mm_kwargs = {}
-
-        # Positions (mrope or standard)
-        if self.uses_mrope:
-            positions = self.mrope_positions[:, :num_input_tokens]
-        else:
-            positions = self.positions[:num_input_tokens]
-
-        # Intermediate tensors sync for PP (if any)
-        if get_pp_group().is_first_rank:
-            intermediate_tensors = None
-        else:
-            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True)
-
-        # Set forward context mainly for resource management and kv connector
-        skip_cuda_graphs = True  # diffusion path does not rely on cuda graphs here
-        with set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_input_tokens,
-                num_tokens_across_dp=num_tokens_across_dp,
-                skip_cuda_graphs=skip_cuda_graphs,
-        ), self.maybe_get_kv_connector_output(
-                scheduler_output) as kv_connector_output:
-
-            if not get_pp_group().is_last_rank:
-                # For non-last PP stages, pass through intermediate tensors.
-                assert intermediate_tensors is not None
-                intermediate_tensors.kv_connector_output = kv_connector_output
-                return intermediate_tensors
-
-            outputs = self._run_diffusion(
-                input_ids=input_ids,
-                positions=positions,
-                intermediate_tensors=intermediate_tensors,
-                inputs_embeds=inputs_embeds,
-                multimodal_kwargs=model_mm_kwargs,
-                logits_indices=logits_indices,
-            )
-        _, multimodal_outputs = (
-            self.extract_multimodal_outputs(outputs))
-
-        # Ensure one tensor per request, map to CPU for output struct
-        pooler_output: List[Optional[torch.Tensor]] = []
-        if isinstance(multimodal_outputs, torch.Tensor):
-            # If model returned a single stacked tensor, split by requests
-            assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
-            for i in range(self.input_batch.num_reqs):
-                pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous())
-        elif isinstance(multimodal_outputs, list):
-            for out in multimodal_outputs:
-                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
-        elif isinstance(multimodal_outputs, dict):
-            for out in multimodal_outputs.values():
-                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
-        else:
-            raise RuntimeError("Unsupported diffusion output type")
-
-        self.eplb_step()
-
-        return OmniModelRunnerOutput(
-            req_ids=self.input_batch.req_ids,
-            req_id_to_index=self.input_batch.req_id_to_index,
-            sampled_token_ids=[],
-            spec_token_ids=None,
-            logprobs=None,
-            prompt_logprobs_dict={},
-            pooler_output=pooler_output,
-            kv_connector_output=kv_connector_output,
-            num_nans_in_logits={},
-        )
-
-    def _run_diffusion(
-        self,
-        *,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        intermediate_tensors: Optional[IntermediateTensors],
-        inputs_embeds: Optional[torch.Tensor],
-        multimodal_kwargs: dict,
-        logits_indices: torch.Tensor,
-    ) -> Union[torch.Tensor, list[torch.Tensor]]:
-        """Runs the diffusion process and returns per-request tensors.
-
-        Tries model interfaces in the following order for maximal compatibility:
-        1) model.sample(condition=..., **kwargs)
-        2) model.forward(condition=..., **kwargs)
-        3) model.diffuse(condition=..., **kwargs)
-        """
-        # Keep inputs identical to AR runner
-        kwargs = dict(
-            input_ids=input_ids,
-            positions=positions,
-            intermediate_tensors=intermediate_tensors,
-            inputs_embeds=inputs_embeds,
-            **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
-            sampling_metadata=self.input_batch.sampling_metadata,
-            logits_index=logits_indices,
-            sampler=self.sampler,
-        )
-
-        # For Qwen 2.5 Omni's current implementation, we only support the forward method
-        if hasattr(self.model, "forward"):
-            return self.model.forward(**kwargs)
-        #TODO: add the diffuse method for other models
-
-        raise RuntimeError(
-            "The loaded model does not expose diffusion interfaces 'sample', "
-            "'forward', or 'diffuse'. Please implement one of them or adapt the runner.")
-
-
-    @torch.inference_mode()
-    def _dummy_run(
-        self,
-        num_tokens: int,
-        capture_attn_cudagraph: bool = False,
-        skip_eplb: bool = False,
-        is_profile: bool = False,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        # Padding for DP
-        num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
-        num_tokens += num_pad
-
-        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
-        # for dummy run with LoRA so that the num_reqs collectively
-        # has num_tokens in total.
-        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
-        max_num_reqs = self.scheduler_config.max_num_seqs
-        num_reqs = min(num_tokens, max_num_reqs)
-        min_tokens_per_req = num_tokens // num_reqs
-        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
-        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
-        assert sum(num_scheduled_tokens_list) == num_tokens
-        assert len(num_scheduled_tokens_list) == num_reqs
-        num_scheduled_tokens = np.array(num_scheduled_tokens_list,
-                                        dtype=np.int32)
-
-        attn_metadata: Optional[dict[str, dict]] = None
-        if capture_attn_cudagraph:
-            attn_metadata = {}
-
-            # Make sure max_model_len is used at the graph capture time.
-            self.seq_lens_np[:num_reqs] = self.max_model_len
-            self.seq_lens_np[num_reqs:] = 0
-            self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
-                                           non_blocking=True)
-
-            for kv_cache_group_id, kv_cache_group_spec in enumerate(
-                    self.kv_cache_config.kv_cache_groups):
-                common_attn_metadata = CommonAttentionMetadata(
-                    query_start_loc=self.query_start_loc[:num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],
-                    seq_lens=self.seq_lens[:num_reqs],
-                    seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.
-                    num_computed_tokens_cpu_tensor[:num_reqs],
-                    num_reqs=num_reqs,
-                    num_actual_tokens=num_tokens,
-                    max_query_len=num_tokens,
-                    block_table_tensor=self.input_batch.block_table[
-                        kv_cache_group_id].get_device_tensor()[:num_reqs],
-                    slot_mapping=self.input_batch.block_table[
-                        kv_cache_group_id].slot_mapping[:num_tokens],
-                    causal=True)
-
-                for attn_group in self.attn_groups[kv_cache_group_id]:
-                    attn_metadata_i = attn_group.metadata_builder \
-                        .build_for_cudagraph_capture(common_attn_metadata)
-                    for layer_name in kv_cache_group_spec.layer_names:
-                        attn_metadata[layer_name] = attn_metadata_i
-
-        with self.maybe_dummy_run_with_lora(self.lora_config,
-                                            num_scheduled_tokens):
-            if self.is_multimodal_model:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds[:num_tokens]
-                model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)
-            else:
-                input_ids = self.input_ids[:num_tokens]
-                inputs_embeds = None
-                model_mm_kwargs = {}
-
-            if self.uses_mrope:
-                positions = self.mrope_positions[:, :num_tokens]
-            else:
-                positions = self.positions[:num_tokens]
-
-            if get_pp_group().is_first_rank:
-                intermediate_tensors = None
-            else:
-                if self.intermediate_tensors is None:
-                    self.intermediate_tensors = (
-                        self.model.make_empty_intermediate_tensors(
-                            batch_size=self.max_num_tokens,
-                            dtype=self.model_config.dtype,
-                            device=self.device))
-
-                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                    num_tokens, None, False)
-
-            # Diffusion path: avoid CUDA graphs; we only use context for resource wiring
-            with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                    attn_metadata,
-                    self.vllm_config,
-                    num_tokens=num_tokens,
-                    num_tokens_across_dp=num_tokens_across_dp):
-                outputs = self.model(
-                    input_ids=input_ids,
-                    positions=positions,
-                    intermediate_tensors=intermediate_tensors,
-                    inputs_embeds=inputs_embeds,
-                    **MultiModalKwargs.as_kwargs(
-                        model_mm_kwargs,
-                        device=self.device,
-                    ),
-                    sampler=None,
-                )
-
-            if self.use_aux_hidden_state_outputs:
-                hidden_states, _ = outputs
-            else:
-                hidden_states = outputs
-
-            # Extract multimodal outputs if present; we ignore them here because
-            # dummy run returns tensors only. The actual diffusion runner returns
-            # multimodal outputs via pooler_output in execute_model.
-            text_hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
-
-            if not skip_eplb:
-                self.eplb_step(is_dummy=True, is_profile=is_profile)
-
-        logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        return text_hidden_states, None
-    
-    @torch.inference_mode()
-    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
-        logger.warning("Dummy sampler run is not implemented for diffusion model")
-        return None
-
-    def profile_run(self) -> None:
-        # Profile with multimodal encoder & encoder cache, similar to base but
-        # without any logits/sampler warming.
-        if self.is_multimodal_model:
-            mm_budget = self.mm_budget
-            assert mm_budget is not None
-
-            # TODO: handle encoder-decoder models once supported.
-            if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
-                (
-                    dummy_modality,
-                    max_tokens,
-                ) = mm_budget.get_modality_with_max_tokens()
-                (
-                    max_mm_items_per_prompt,
-                    max_mm_items_per_batch,
-                ) = mm_budget.get_max_items(dummy_modality, max_tokens)
-
-                batched_dummy_mm_inputs = self._get_mm_dummy_batch(
-                    dummy_modality,
-                    max_mm_items_per_batch,
-                )
-
-                dummy_encoder_outputs = self.model.get_multimodal_embeddings(
-                    **batched_dummy_mm_inputs)
-
-                sanity_check_mm_encoder_outputs(
-                    dummy_encoder_outputs,
-                    expected_num_items=max_mm_items_per_batch,
-                )
-
-                self.encoder_cache["tmp"] = dict(
-                    enumerate(dummy_encoder_outputs))
-
-        hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
-        if get_pp_group().is_last_rank:
-            pass  # No sampler/pooler warmup for diffusion
-        self._sync_device()
-        del hidden_states
-        self.encoder_cache.clear()
-        gc.collect()

vllm_omni/worker/gpu_model_runner.py

@@ -1,755 +0,0 @@
-
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
-
-import numpy as np
-import torch
-
-import vllm.envs as envs
-from vllm.distributed.parallel_state import get_pp_group, get_tp_group
-from vllm.forward_context import set_forward_context
-from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces_base import VllmModelForPooling
-from vllm.multimodal.inputs import MultiModalKwargs
-from vllm.sequence import IntermediateTensors
-from vllm.utils import LazyLoader
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
-from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
-from vllm.v1.worker.gpu_input_batch import CachedRequestState
-from vllm.sampling_params import SamplingType
-from vllm.distributed.kv_transfer import has_kv_transfer_group
-from vllm.utils import round_up
-from vllm.v1.spec_decode.eagle import EagleProposer
-
-from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-
-from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
-
-if TYPE_CHECKING:
-    from vllm.v1.core.sched.output import SchedulerOutput
-else:
-    xgr = LazyLoader("xgr", globals(), "xgrammar")
-    xgr_torch_compile = LazyLoader(
-        "xgr_torch_compile", globals(),
-        "xgrammar.kernels.apply_token_bitmask_inplace_torch_compile")
-
-logger = init_logger(__name__)
-
-
-class OmniGPUModelRunner(GPUModelRunner):
-    def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
-        """Update the cached states and the persistent batch with the scheduler
-        output.
-
-        The updated states are used by the `_prepare_inputs` function to create
-        the input GPU tensors for the model.
-
-        The SamplingMetadata is updated and copied to the GPU if there is a
-        new/resumed/paused/finished request in the batch.
-        """
-        # Remove finished requests from the cached states.
-        for req_id in scheduler_output.finished_req_ids:
-            self.requests.pop(req_id, None)
-            self.encoder_cache.pop(req_id, None)
-        # Remove the finished requests from the persistent batch.
-        # NOTE(woosuk): There could be an edge case where finished_req_ids and
-        # scheduled_req_ids overlap. This happens when a request is aborted and
-        # then resubmitted with the same ID. In this case, we treat them as two
-        # distinct requests - clearing the cached states for the first request
-        # and handling the second as a new request.
-        for req_id in scheduler_output.finished_req_ids:
-            self.input_batch.remove_request(req_id)
-
-        # Free the cached encoder outputs.
-        for req_id, input_id in scheduler_output.free_encoder_input_ids:
-            encoder_outputs = self.encoder_cache.get(req_id)
-            if encoder_outputs is not None:
-                encoder_outputs.pop(input_id, None)
-                if not encoder_outputs:
-                    self.encoder_cache.pop(req_id, None)
-
-        # Remove the unscheduled requests from the persistent batch.
-        # NOTE(woosuk): The unscheduled requests are either preempted requests
-        # or running requests that are not scheduled in this step. We remove
-        # them from the persistent batch but keep their cached states since
-        # they will be scheduled again sometime in the future.
-        scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
-        cached_req_ids = self.input_batch.req_id_to_index.keys()
-        unscheduled_req_ids = cached_req_ids - scheduled_req_ids
-        # NOTE(woosuk): The persistent batch optimization assumes that
-        # consecutive batches contain mostly the same requests. If batches
-        # have low request overlap (e.g., alternating between two distinct
-        # sets of requests), this optimization becomes very inefficient.
-        for req_id in unscheduled_req_ids:
-            self.input_batch.remove_request(req_id)
-
-        req_ids_to_add: list[str] = []
-        # Add new requests to the cached states.
-        for new_req_data in scheduler_output.scheduled_new_reqs:
-            req_id = new_req_data.req_id
-            sampling_params = new_req_data.sampling_params
-            pooling_params = new_req_data.pooling_params
-
-            if sampling_params and \
-                sampling_params.sampling_type == SamplingType.RANDOM_SEED:
-                generator = torch.Generator(device=self.device)
-                generator.manual_seed(sampling_params.seed)
-            else:
-                generator = None
-
-            if pooling_params:
-                assert (task := pooling_params.task) is not None, (
-                    "You did not set `task` in the API")
-
-                model = cast(VllmModelForPooling, self.model)
-                to_update = model.pooler.get_pooling_updates(task)
-                to_update.apply(pooling_params)
-
-            self.requests[req_id] = CachedRequestState(
-                req_id=req_id,
-                prompt_token_ids=new_req_data.prompt_token_ids,
-                mm_inputs=new_req_data.mm_inputs,
-                mm_positions=new_req_data.mm_positions,
-                sampling_params=sampling_params,
-                pooling_params=pooling_params,
-                generator=generator,
-                block_ids=new_req_data.block_ids,
-                num_computed_tokens=new_req_data.num_computed_tokens,
-                output_token_ids=[],
-                lora_request=new_req_data.lora_request,
-            )
-
-            # If prompt embeddings are provided, decode and attach to inter_data
-            try:
-                if getattr(new_req_data, "prompt_embeds", None) is not None:
-                    payload = new_req_data.prompt_embeds
-                    dtype = getattr(np, payload.dtype)
-                    arr = np.frombuffer(payload.data, dtype=dtype)
-                    arr = arr.reshape(payload.shape)
-                    pe_cpu = torch.from_numpy(arr)
-                    # Store temporarily on CPU; later moved to device in builder
-                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
-                    # Also replace payload with Tensor for user visibility in scheduler_output
-                    try:
-                        new_req_data.prompt_embeds = pe_cpu  # type: ignore[assignment]
-                    except Exception:
-                        pass
-            except Exception:
-                pass
-            # Decode additional_information payloads (dictionary)
-            try:
-                if getattr(new_req_data, "additional_information", None) is not None:
-                    payload_info = new_req_data.additional_information
-                    info_dict = {}
-                    if isinstance(payload_info, dict):
-                        info_dict = payload_info
-                    else:
-                        from vllm.v1.engine import AdditionalInformationPayload
-                        if isinstance(payload_info, AdditionalInformationPayload):
-                            for k, entry in payload_info.entries.items():
-                                if entry.tensor_data is not None:
-                                    dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
-                                    arr = np.frombuffer(entry.tensor_data, dtype=dt)
-                                    arr = arr.reshape(entry.tensor_shape)
-                                    info_dict[k] = torch.from_numpy(arr)
-                                else:
-                                    info_dict[k] = entry.list_data
-                    if info_dict:
-                        setattr(self.requests[req_id], "additional_information_cpu", info_dict)
-            except Exception:
-                pass
-
-            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
-            if self.uses_mrope:
-                image_grid_thw = []
-                video_grid_thw = []
-                second_per_grid_ts = []
-                audio_feature_lengths = []
-                use_audio_in_video = False
-                for mm_input in self.requests[req_id].mm_inputs:
-                    if mm_input.get("image_grid_thw") is not None:
-                        image_grid_thw.extend(
-                            mm_input["image_grid_thw"].tolist())
-                    if mm_input.get("video_grid_thw") is not None:
-                        video_grid_thw.extend(
-                            mm_input["video_grid_thw"].tolist())
-                    if mm_input.get("second_per_grid_ts") is not None:
-                        second_per_grid_ts.extend(
-                            mm_input["second_per_grid_ts"])
-                    if mm_input.get("audio_feature_lengths") is not None:
-                        audio_feature_lengths.extend(
-                            mm_input["audio_feature_lengths"])
-                    if mm_input.get("use_audio_in_video") is True:
-                        use_audio_in_video = True
-
-                hf_config = self.model_config.hf_config
-
-                self.requests[req_id].mrope_positions, \
-                    self.requests[req_id].mrope_position_delta = \
-                    MRotaryEmbedding.get_input_positions_tensor(
-                        self.requests[req_id].prompt_token_ids,
-                        hf_config=hf_config,
-                        image_grid_thw=image_grid_thw,
-                        video_grid_thw=video_grid_thw,
-                        second_per_grid_ts=second_per_grid_ts,
-                        audio_feature_lengths=audio_feature_lengths,
-                        use_audio_in_video=use_audio_in_video,
-                    )
-
-            req_ids_to_add.append(req_id)
-
-        # Update the states of the running/resumed requests.
-        is_last_rank = get_pp_group().is_last_rank
-        req_data = scheduler_output.scheduled_cached_reqs
-        for i, req_id in enumerate(req_data.req_ids):
-            req_state = self.requests[req_id]
-            num_computed_tokens = req_data.num_computed_tokens[i]
-            new_block_ids = req_data.new_block_ids[i]
-            resumed_from_preemption = req_data.resumed_from_preemption[i]
-
-            # Update the cached states.
-            req_state.num_computed_tokens = num_computed_tokens
-
-            if not is_last_rank:
-                # When using PP, the scheduler sends the sampled tokens back,
-                # because there's no direct communication between the first-
-                # stage worker and the last-stage worker.
-                new_token_ids = req_data.new_token_ids[i]
-                # Add the sampled token(s) from the previous step (if any).
-                # This doesn't include "unverified" tokens like spec tokens.
-                num_new_tokens = (num_computed_tokens + len(new_token_ids) -
-                                  req_state.num_tokens)
-                if num_new_tokens == 1:
-                    # Avoid slicing list in most common case.
-                    req_state.output_token_ids.append(new_token_ids[-1])
-                elif num_new_tokens > 0:
-                    req_state.output_token_ids.extend(
-                        new_token_ids[-num_new_tokens:])
-
-            # Update the block IDs.
-            if not resumed_from_preemption:
-                # Append the new blocks to the existing block IDs.
-                for block_ids, new_ids in zip(req_state.block_ids,
-                                              new_block_ids):
-                    block_ids.extend(new_ids)
-            else:
-                # The request is resumed from preemption.
-                # Replace the existing block IDs with the new ones.
-                req_state.block_ids = new_block_ids
-
-            req_index = self.input_batch.req_id_to_index.get(req_id)
-            if req_index is None:
-                # The request is not in the persistent batch.
-                # The request was either preempted and resumed later, or was not
-                # scheduled in the previous step and needs to be added again.
-                req_ids_to_add.append(req_id)
-                continue
-
-            # Update the persistent batch.
-            self.input_batch.num_computed_tokens_cpu[req_index] = (
-                num_computed_tokens)
-            self.input_batch.block_table.append_row(new_block_ids, req_index)
-
-            # For the last rank, we don't need to update the token_ids_cpu
-            # because the sampled tokens are already cached.
-            if not is_last_rank:
-                # Add new_token_ids to token_ids_cpu.
-                start_token_index = num_computed_tokens
-                end_token_index = num_computed_tokens + len(new_token_ids)
-                self.input_batch.token_ids_cpu[
-                    req_index,
-                    start_token_index:end_token_index] = new_token_ids
-                self.input_batch.num_tokens_no_spec[
-                    req_index] = end_token_index
-                self.input_batch.num_tokens[req_index] = end_token_index
-
-            # Add spec_token_ids to token_ids_cpu.
-            spec_token_ids = (
-                scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()))
-            if spec_token_ids:
-                num_spec_tokens = len(spec_token_ids)
-                start_index = self.input_batch.num_tokens_no_spec[req_index]
-                end_token_index = start_index + num_spec_tokens
-                self.input_batch.token_ids_cpu[
-                    req_index, start_index:end_token_index] = spec_token_ids
-                # NOTE(woosuk): `num_tokens` here may include spec tokens.
-                self.input_batch.num_tokens[req_index] += num_spec_tokens
-
-        # Add the new or resumed requests to the persistent batch.
-        # The smaller empty indices are filled first.
-        for req_id in req_ids_to_add:
-            req_state = self.requests[req_id]
-            self.input_batch.add_request(req_state)
-
-        # Condense the batched states if there are gaps left by removed requests
-        self.input_batch.condense()
-        # Allow attention backend to reorder the batch, potentially
-        self._may_reorder_batch(scheduler_output)
-        # Refresh batch metadata with any pending updates.
-        self.input_batch.refresh_metadata()
-    
-    @torch.inference_mode()
-    def execute_model(
-        self,
-        scheduler_output: "SchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-    ) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
-        self._update_states(scheduler_output)
-        if not scheduler_output.total_num_scheduled_tokens:
-            if not has_kv_transfer_group():
-                # Return empty ModelRunnerOutput if there's no work to do.
-                return EMPTY_MODEL_RUNNER_OUTPUT
-
-            return self.kv_connector_no_forward(scheduler_output,
-                                                self.vllm_config)
-
-        # Prepare the decoder inputs.
-        (attn_metadata, attention_cuda_graphs, logits_indices,
-         spec_decode_metadata, num_scheduled_tokens_np,
-         spec_decode_common_attn_metadata) = (
-             self._prepare_inputs(scheduler_output))
-
-        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        if (self.use_cuda_graph
-                and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
-            # Use piecewise CUDA graphs.
-            # Add padding to the batch size.
-            num_input_tokens = self.vllm_config.pad_for_cudagraph(
-                num_scheduled_tokens)
-        else:
-            # Eager mode.
-            # Pad tokens to multiple of tensor_parallel_size when
-            # enabled collective fusion for SP
-            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
-            if self.compilation_config.pass_config. \
-                enable_sequence_parallelism and tp_size > 1:
-                num_input_tokens = round_up(num_scheduled_tokens, tp_size)
-            else:
-                num_input_tokens = num_scheduled_tokens
-
-        # Padding for DP
-        num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
-        num_input_tokens += num_pad
-
-        # _prepare_inputs may reorder the batch, so we must gather multi
-        # modal outputs after that to ensure the correct order
-        if self.is_multimodal_model:
-            # Run the multimodal encoder if any.
-            self._execute_mm_encoder(scheduler_output)
-            mm_embeds = self._gather_mm_embeddings(scheduler_output)
-        else:
-            mm_embeds = []
-
-        if get_pp_group().is_first_rank:
-            # Always prepare inputs_embeds on the first PP rank so both
-            # multimodal and text-only models can consume embeddings.
-            inputs_embeds_scheduled = self.model.get_input_embeddings(
-                input_ids=self.input_ids[:num_scheduled_tokens],
-                multimodal_embeddings=(mm_embeds or None)
-                if self.is_multimodal_model else None,
-            )
-
-            # TODO(woosuk): Avoid the copy. Optimize.
-            self.inputs_embeds[:num_scheduled_tokens].copy_(
-                inputs_embeds_scheduled)
-
-            # Overlay per-request custom prompt_embeds for the prompt portion
-            for req_index, req_id in enumerate(self.input_batch.req_ids):
-                req_state = self.requests[req_id]
-                pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
-                if pe_cpu is None:
-                    continue
-                num_computed_tokens = int(
-                    self.input_batch.num_computed_tokens_cpu[req_index])
-                prompt_len = len(req_state.prompt_token_ids)
-                prompt_remaining = max(0, prompt_len - num_computed_tokens)
-                sched_tokens = int(num_scheduled_tokens_np[req_index])
-                overlay_len = min(sched_tokens, prompt_remaining)
-                if overlay_len <= 0:
-                    continue
-                src = pe_cpu[num_computed_tokens:
-                             num_computed_tokens + overlay_len].to(
-                                 dtype=self.dtype,
-                                 device=self.device,
-                                 non_blocking=True)
-                start_offset = int(self.query_start_loc_cpu[req_index])
-                self.inputs_embeds[start_offset:start_offset + overlay_len]\
-                    .copy_(src)
-
-            input_ids = self.input_ids[:num_input_tokens]  # preserved for APIs
-            inputs_embeds = self.inputs_embeds[:num_input_tokens]
-            model_mm_kwargs = (self._extract_mm_kwargs(scheduler_output)
-                               if self.is_multimodal_model else {})
-        else:
-            # For non-first PP ranks, use token ids as usual; embeddings are
-            # only consumed on the first rank.
-            input_ids = self.input_ids[:num_input_tokens]
-            inputs_embeds = None
-            model_mm_kwargs = {}
-        if self.uses_mrope:
-            positions = self.mrope_positions[:, :num_input_tokens]
-        else:
-            positions = self.positions[:num_input_tokens]
-
-        if get_pp_group().is_first_rank:
-            intermediate_tensors = None
-        else:
-            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True)
-
-        # Some attention backends only support CUDA Graphs in pure decode.
-        # If attention doesn't support CUDA Graphs for this batch, but we
-        # compiled with full CUDA graphs, we have to skip them entirely.
-        skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
-
-        # Run the model.
-        # Use persistent buffers for CUDA graphs.
-        with set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_input_tokens,
-                num_tokens_across_dp=num_tokens_across_dp,
-                skip_cuda_graphs=skip_cuda_graphs,
-        ), self.maybe_get_kv_connector_output(
-                scheduler_output) as kv_connector_output:
-
-            model_output = self.model(
-                input_ids=input_ids,
-                positions=positions,
-                intermediate_tensors=intermediate_tensors,
-                inputs_embeds=inputs_embeds,
-                **MultiModalKwargs.as_kwargs(
-                    model_mm_kwargs,
-                    device=self.device,
-                ),
-                sampling_metadata=self.input_batch.sampling_metadata,
-                logits_index=logits_indices,
-                sampler=self.sampler,
-            )
-
-        if self.use_aux_hidden_state_outputs:
-            hidden_states, aux_hidden_states = model_output
-        else:
-            hidden_states = model_output
-            aux_hidden_states = None
-
-        text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
-
-        # Broadcast PP output for external_launcher (torchrun)
-        # to make sure we are synced across pp ranks
-        # TODO: Support overlapping mirco-batches
-        # https://github.com/vllm-project/vllm/issues/18019
-        broadcast_pp_output = \
-            self.parallel_config.distributed_executor_backend \
-            == "external_launcher" and len(get_pp_group().ranks) > 0
-        if not get_pp_group().is_last_rank:
-            # For mid-pipeline stages, return the hidden states.
-            assert isinstance(text_hidden_states, IntermediateTensors)
-            if not broadcast_pp_output:
-                text_hidden_states.kv_connector_output = kv_connector_output
-                return text_hidden_states
-            get_pp_group().send_tensor_dict(text_hidden_states.tensors,
-                                            all_gather_group=get_tp_group())
-            logits = None
-        else:
-            if self.input_batch.pooling_params:
-                return self._pool(text_hidden_states, num_scheduled_tokens,
-                                  num_scheduled_tokens_np, kv_connector_output)
-
-            sample_hidden_states = text_hidden_states[logits_indices]
-            logits = self.model.compute_logits(sample_hidden_states, None)
-        if broadcast_pp_output:
-            model_output_broadcast_data = {
-                "logits": logits.contiguous(),
-            } if logits is not None else {}
-            model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
-                model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)
-            assert model_output_broadcast_data is not None
-            logits = model_output_broadcast_data["logits"]
-
-        # Apply structured output bitmasks if present
-        if scheduler_output.grammar_bitmask is not None:
-            self.apply_grammar_bitmask(scheduler_output, logits)
-
-        # Sample the next token and get logprobs if needed.
-        sampling_metadata = self.input_batch.sampling_metadata
-        if spec_decode_metadata is None:
-            sampler_output = self.sampler(
-                logits=logits,
-                sampling_metadata=sampling_metadata,
-            )
-        else:
-            # When indexing with a tensor (bonus_logits_indices), PyTorch
-            # creates a new tensor with separate storage from the original
-            # logits tensor. This means any in-place operations on bonus_logits
-            # won't affect the original logits tensor.
-            assert logits is not None
-            bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
-            sampler_output = self.sampler(
-                logits=bonus_logits,
-                sampling_metadata=sampling_metadata,
-            )
-            bonus_token_ids = sampler_output.sampled_token_ids
-
-            # Just like `bonus_logits`, `target_logits` is a new tensor with
-            # separate storage from the original `logits` tensor. Therefore,
-            # it is safe to update `target_logits` in place.
-            target_logits = logits[spec_decode_metadata.target_logits_indices]
-            output_token_ids = self.rejection_sampler(
-                spec_decode_metadata,
-                None,  # draft_probs
-                target_logits,
-                bonus_token_ids,
-                sampling_metadata,
-            )
-            sampler_output.sampled_token_ids = output_token_ids
-
-        num_nans_in_logits = {}
-        if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
-            num_nans_in_logits = self._get_nans_in_logits(logits)
-
-        # TODO(woosuk): The following loop can be slow since it iterates over
-        # the requests one by one. Optimize.
-        discard_sampled_tokens_req_indices = []
-        for i, req_id in enumerate(self.input_batch.req_ids):
-            req_state = self.requests[req_id]
-            seq_len = (req_state.num_computed_tokens +
-                       scheduler_output.num_scheduled_tokens[req_id])
-            if seq_len < req_state.num_tokens:
-                # Ignore the sampled token for partial prefills.
-                # Rewind the generator state as if the token was not sampled.
-                # This relies on cuda-specific torch-internal impl details
-                generator = self.input_batch.generators.get(i)
-                if generator is not None:
-                    generator.set_offset(generator.get_offset() - 4)
-                # Record the index of the request that should not be sampled,
-                # so that we could clear the sampled tokens before returning.
-                discard_sampled_tokens_req_indices.append(i)
-
-        # NOTE: GPU -> CPU Sync happens here.
-        # Move as many CPU operations as possible before this sync point.
-        logprobs_tensors = sampler_output.logprobs_tensors
-        logprobs_lists = logprobs_tensors.tolists() \
-            if logprobs_tensors is not None else None
-
-        # Compute prompt logprobs if needed.
-        prompt_logprobs_dict = self._get_prompt_logprobs_dict(
-            text_hidden_states[:num_scheduled_tokens],
-            scheduler_output,
-        )
-
-        # Get the valid generated tokens.
-        import os
-        sampled_token_ids = sampler_output.sampled_token_ids
-        max_gen_len = sampled_token_ids.shape[-1]
-        if max_gen_len == 1:
-            # No spec decode tokens.
-            valid_sampled_token_ids = sampled_token_ids.tolist()
-        else:
-            # Includes spec decode tokens.
-            valid_sampled_token_ids = self.rejection_sampler.parse_output(
-                sampled_token_ids,
-                self.input_batch.vocab_size,
-            )
-        # Mask out the sampled tokens that should not be sampled.
-        for i in discard_sampled_tokens_req_indices:
-            valid_sampled_token_ids[i].clear()
-
-        # Cache the sampled tokens in the model runner, so that the scheduler
-        # doesn't need to send them back.
-        # NOTE(woosuk): As an exception, when using PP, the scheduler sends
-        # the sampled tokens back, because there's no direct communication
-        # between the first-stage worker and the last-stage worker.
-        for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
-            if not sampled_ids:
-                continue
-
-            start_idx = self.input_batch.num_tokens_no_spec[req_idx]
-            end_idx = start_idx + len(sampled_ids)
-            assert end_idx <= self.max_model_len, (
-                "Sampled token IDs exceed the max model length. "
-                f"Total number of tokens: {end_idx} > max_model_len: "
-                f"{self.max_model_len}")
-
-            self.input_batch.token_ids_cpu[req_idx,
-                                           start_idx:end_idx] = sampled_ids
-            self.input_batch.num_tokens_no_spec[req_idx] = end_idx
-            self.input_batch.num_tokens[req_idx] = end_idx
-            req_id = self.input_batch.req_ids[req_idx]
-            req_state = self.requests[req_id]
-            req_state.output_token_ids.extend(sampled_ids)
-
-        if not self.speculative_config:
-            # Speculative decoding is not enabled.
-            spec_token_ids = None
-        else:
-            assert spec_decode_common_attn_metadata is not None
-            spec_token_ids = self.propose_draft_token_ids(
-                scheduler_output,
-                valid_sampled_token_ids,
-                sampling_metadata,
-                text_hidden_states,
-                sample_hidden_states,
-                aux_hidden_states,
-                spec_decode_metadata,
-                spec_decode_common_attn_metadata,
-            )
-
-        self.eplb_step()
-
-        return OmniModelRunnerOutput(
-            req_ids=self.input_batch.req_ids,
-            req_id_to_index=self.input_batch.req_id_to_index,
-            sampled_token_ids=valid_sampled_token_ids,
-            spec_token_ids=spec_token_ids,
-            logprobs=logprobs_lists,
-            prompt_logprobs_dict=prompt_logprobs_dict,
-            pooler_output=[],
-            kv_connector_output=kv_connector_output,
-            num_nans_in_logits=num_nans_in_logits,
-            multimodal_outputs=multimodal_outputs,
-        )
-    
-    @torch.inference_mode()
-    def extract_multimodal_outputs(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, dict]:
-        if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs:
-            text_hidden_states = hidden_states.text_hidden_states
-            multimodal_outputs = hidden_states.multimodal_outputs
-
-        else:
-            text_hidden_states = hidden_states
-            multimodal_outputs = {}
-        return text_hidden_states, multimodal_outputs
-    
-    @torch.inference_mode()
-    def _dummy_run(
-        self,
-        num_tokens: int,
-        capture_attn_cudagraph: bool = False,
-        skip_eplb: bool = False,
-        is_profile: bool = False,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-
-        # Padding for DP
-        num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
-        num_tokens += num_pad
-
-        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
-        # for dummy run with LoRA so that the num_reqs collectively
-        # has num_tokens in total.
-        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
-        max_num_reqs = self.scheduler_config.max_num_seqs
-        num_reqs = min(num_tokens, max_num_reqs)
-        min_tokens_per_req = num_tokens // num_reqs
-        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
-        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
-        assert sum(num_scheduled_tokens_list) == num_tokens
-        assert len(num_scheduled_tokens_list) == num_reqs
-        num_scheduled_tokens = np.array(num_scheduled_tokens_list,
-                                        dtype=np.int32)
-
-        attn_metadata: Optional[dict[str, Any]] = None
-        if capture_attn_cudagraph:
-            attn_metadata = {}
-
-            # Make sure max_model_len is used at the graph capture time.
-            self.seq_lens_np[:num_reqs] = self.max_model_len
-            self.seq_lens_np[num_reqs:] = 0
-            self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
-                                           non_blocking=True)
-
-            for kv_cache_group_id, kv_cache_group_spec in enumerate(
-                    self.kv_cache_config.kv_cache_groups):
-                common_attn_metadata = CommonAttentionMetadata(
-                    query_start_loc=self.query_start_loc[:num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs +
-                                                                 1],
-                    seq_lens=self.seq_lens[:num_reqs],
-                    seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.
-                    num_computed_tokens_cpu_tensor[:num_reqs],
-                    num_reqs=num_reqs,
-                    num_actual_tokens=num_tokens,
-                    max_query_len=num_tokens,
-                    block_table_tensor=self.input_batch.block_table[
-                        kv_cache_group_id].get_device_tensor()[:num_reqs],
-                    slot_mapping=self.input_batch.
-                    block_table[kv_cache_group_id].slot_mapping[:num_tokens],
-                    causal=True)
-
-                for attn_group in self.attn_groups[kv_cache_group_id]:
-                    attn_metadata_i = attn_group.metadata_builder\
-                        .build_for_cudagraph_capture(common_attn_metadata)
-                    for layer_name in kv_cache_group_spec.layer_names:
-                        attn_metadata[layer_name] = attn_metadata_i
-
-        with self.maybe_dummy_run_with_lora(self.lora_config,
-                                            num_scheduled_tokens):
-            if self.is_multimodal_model:
-                input_ids = None
-                inputs_embeds = self.inputs_embeds[:num_tokens]
-                model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)
-            else:
-                input_ids = self.input_ids[:num_tokens]
-                inputs_embeds = None
-                model_mm_kwargs = {}
-
-            if self.uses_mrope:
-                positions = self.mrope_positions[:, :num_tokens]
-            else:
-                positions = self.positions[:num_tokens]
-
-            if get_pp_group().is_first_rank:
-                intermediate_tensors = None
-            else:
-                if self.intermediate_tensors is None:
-                    self.intermediate_tensors = (
-                        self.model.make_empty_intermediate_tensors(
-                            batch_size=self.max_num_tokens,
-                            dtype=self.model_config.dtype,
-                            device=self.device))
-
-                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                    num_tokens, None, False)
-
-            with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                    attn_metadata,
-                    self.vllm_config,
-                    num_tokens=num_tokens,
-                    num_tokens_across_dp=num_tokens_across_dp):
-                outputs = self.model(
-                    input_ids=input_ids,
-                    positions=positions,
-                    intermediate_tensors=intermediate_tensors,
-                    inputs_embeds=inputs_embeds,
-                    **MultiModalKwargs.as_kwargs(
-                        model_mm_kwargs,
-                        device=self.device,
-                    ),
-                    sampler = None
-                )
-
-            if self.use_aux_hidden_state_outputs:
-                hidden_states, _ = outputs
-            else:
-                hidden_states = outputs
-
-            text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
-
-
-            if self.speculative_config and self.speculative_config.use_eagle():
-                assert isinstance(self.drafter, EagleProposer)
-                self.drafter.dummy_run(num_tokens)
-
-        # This is necessary to avoid blocking DP.
-        # For dummy runs, we typically skip EPLB since we don't have any real
-        # requests to process.
-        # However, in DP settings, there may be cases when some DP ranks do
-        # not have any requests to process, so they're executing dummy batches.
-        # In such cases, we still have to trigger EPLB to make sure
-        # ranks execute the rearrangement in synchronization.
-        if not skip_eplb:
-            self.eplb_step(is_dummy=True, is_profile=is_profile)
-
-        logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        return text_hidden_states, text_hidden_states[logit_indices]

vllm_omni/worker/model_runner.py

@@ -1,199 +0,0 @@
-from dataclasses import dataclass
-from typing import Optional, Dict, Any, List
-import itertools
-
-import torch
-
-from vllm.worker.model_runner import ModelInputForGPU, ModelInputForGPUBuilder
-
-from vllm.utils import flatten_2d_lists, async_tensor_h2d
-from vllm.lora.layers import LoRAMapping
-from vllm.sequence import MultiModalKwargs
-
-
-@dataclass(frozen=True)
-class OmniModelInputForGPU(ModelInputForGPU):
-    # New: optional additional information dict for current scheduled tokens
-    additional_information: Optional[Dict[str, Any]] = None
-
-
-class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder):
-    def build(self) -> OmniModelInputForGPU:
-        """Finalize the builder intermediate data and
-        create on-device tensors.
-        """
-        # Combine and flatten intermediate data.
-        input_tokens = list[int]()
-        inputs_embeds_list = list[torch.Tensor]()
-        token_types = list[int]()
-        for inter_data in self.inter_data_list:
-            for cur_input_tokens in inter_data.input_tokens:
-                input_tokens.extend(cur_input_tokens)
-            for cur_token_types in inter_data.token_types:
-                token_types.extend(cur_token_types)
-            if inter_data.inputs_embeds is not None:
-                inputs_embeds_list.append(
-                    inter_data.inputs_embeds.to(
-                        dtype=self.runner.model_config.dtype,
-                        device=self.runner.device))
-            # Support v1 direct-transfer prompt embeds attached on request state
-            if inter_data.inputs_embeds is None:
-                try:
-                    # Locate req_id -> request state to fetch decoded CPU embeds
-                    req_id = inter_data.request_id  # type: ignore[attr-defined]
-                    req_state = getattr(self.runner, "requests", {}).get(req_id)
-                    pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
-                    if pe_cpu is not None:
-                        inter_data.inputs_embeds = pe_cpu.to(
-                            dtype=self.runner.model_config.dtype,
-                            device=self.runner.device)
-                        inputs_embeds_list.append(inter_data.inputs_embeds)
-                except Exception:
-                    pass
-        inputs_embeds: Optional[torch.Tensor]
-        if len(inputs_embeds_list) == 0:
-            inputs_embeds = None
-        else:
-            inputs_embeds = torch.cat(inputs_embeds_list, dim=0).to(
-                dtype=self.runner.model_config.dtype,
-                device=self.runner.device)
-            assert len(inputs_embeds) == len(input_tokens)
-
-        if not input_tokens and inputs_embeds is None:
-            # This may happen when all prefill requests hit
-            # prefix caching and there is no decode request.
-            return self.model_input_cls()
-
-        mrope_input_positions: Optional[List[List[int]]] = None
-        if any(inter_data.mrope_input_positions is not None
-               for inter_data in self.inter_data_list):
-            mrope_input_positions = [[] for _ in range(3)]
-            for idx in range(3):
-                for inter_data in self.inter_data_list:
-                    msections = inter_data.mrope_input_positions
-                    if msections is None:
-                        for _seq_input_positions in inter_data.input_positions:
-                            mrope_input_positions[idx].extend(
-                                _seq_input_positions)
-                    else:
-                        for _seq_mrope_input_positions in msections:
-                            mrope_input_positions[idx].extend(
-                                _seq_mrope_input_positions[idx])
-            input_positions = None
-        else:
-            input_positions = []
-            for inter_data in self.inter_data_list:
-                for cur_input_positions in inter_data.input_positions:
-                    input_positions.extend(cur_input_positions)
-
-        seq_lens = []
-        query_lens = []
-        max_decode_seq_len = 0
-        max_encoder_seq_len = 0
-        for inter_data in self.inter_data_list:
-            seq_lens.extend(inter_data.seq_lens)
-            query_lens.extend(inter_data.query_lens)
-            if not inter_data.is_prompt:
-                max_decode_seq_len = max(max_decode_seq_len,
-                                         max(inter_data.seq_lens))
-                if self.runner.model_config.is_encoder_decoder:
-                    max_encoder_seq_len = max(max_encoder_seq_len,
-                                              inter_data.encoder_seq_len)
-
-        # Mapping from request IDs to sequence IDs. Used for Jamba models
-        # that manages the cache by itself.
-        request_ids_to_seq_ids = {
-            data.request_id: data.seq_ids
-            for data in self.inter_data_list
-        }
-
-        cuda_graph_pad_size = self._get_cuda_graph_pad_size(
-            num_seqs=len(seq_lens),
-            max_decode_seq_len=max_decode_seq_len,
-            max_encoder_seq_len=max_encoder_seq_len)
-
-        batch_size = len(input_tokens)
-        if cuda_graph_pad_size != -1:
-            # If cuda graph can be used, pad tensors accordingly.
-            # See `capture_model` API for more details.
-            # vLLM uses cuda graph only for decoding requests.
-            batch_size += cuda_graph_pad_size
-
-        # Tokens and positions.
-        if cuda_graph_pad_size:
-            input_tokens.extend(itertools.repeat(0, cuda_graph_pad_size))
-        assert self.runner.device is not None
-        input_tokens_tensor = async_tensor_h2d(input_tokens, torch.long,
-                                               self.runner.device,
-                                               self.runner.pin_memory)
-
-        token_types_tensor = async_tensor_h2d(token_types, torch.long,
-                                               self.runner.device,
-                                               self.runner.pin_memory) \
-                                                if token_types else None
-
-        if mrope_input_positions is not None:
-            for idx in range(3):
-                mrope_input_positions[idx].extend(
-                    itertools.repeat(0, cuda_graph_pad_size))
-            input_positions_tensor = async_tensor_h2d(mrope_input_positions,
-                                                      torch.long,
-                                                      self.runner.device,
-                                                      self.runner.pin_memory)
-        else:
-            input_positions.extend(itertools.repeat(0, cuda_graph_pad_size))
-            input_positions_tensor = async_tensor_h2d(input_positions,
-                                                      torch.long,
-                                                      self.runner.device,
-                                                      self.runner.pin_memory)
-        # Sequence and query lengths.
-        if cuda_graph_pad_size:
-            seq_lens.extend(itertools.repeat(1, cuda_graph_pad_size))
-
-        # Attention metadata.
-        attn_metadata = self.attn_metadata_builder.build(
-            seq_lens, query_lens, cuda_graph_pad_size, batch_size)
-
-        # LoRA data.
-        lora_requests = set()
-        lora_mapping = None
-        if self.enable_lora:
-            lora_requests = set(r for data in self.inter_data_list
-                                for r in data.lora_requests)
-            lora_index_mapping = flatten_2d_lists([
-                flatten_2d_lists(inter_data.lora_index_mapping)
-                for inter_data in self.inter_data_list
-            ])
-            if cuda_graph_pad_size:
-                lora_index_mapping.extend(
-                    itertools.repeat(0, cuda_graph_pad_size))
-            lora_prompt_mapping = flatten_2d_lists([
-                flatten_2d_lists(inter_data.lora_prompt_mapping)
-                for inter_data in self.inter_data_list
-            ])
-
-            lora_mapping = LoRAMapping(
-                **dict(index_mapping=lora_index_mapping,
-                       prompt_mapping=lora_prompt_mapping,
-                       is_prefill=not self.decode_only))
-
-        # Multi-modal data.
-        multi_modal_kwargs_list = [
-            data.multi_modal_kwargs for data in self.inter_data_list
-            if data.multi_modal_kwargs is not None
-        ]
-        multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
-
-        return self.model_input_cls(
-            input_tokens=input_tokens_tensor,
-            inputs_embeds=inputs_embeds,
-            input_positions=input_positions_tensor,
-            token_types=token_types_tensor,
-            attn_metadata=attn_metadata,
-            seq_lens=seq_lens,
-            query_lens=query_lens,
-            lora_mapping=lora_mapping,
-            lora_requests=lora_requests,
-            multi_modal_kwargs=multi_modal_kwargs,
-            request_ids_to_seq_ids=request_ids_to_seq_ids,
-            finished_requests_ids=self.finished_requests_ids)