update

ray + mooncake connector work on 2 nodes, but needes several updates\n1.start ray inside vllm-omni \n2.placement group optimize \n3.fix close

.gitignore

@@ -162,9 +162,6 @@ docker/
 # scripts
 /scripts/
 
-# tests
-tests/
-
 # PyCharm
 #  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
 #  be added to the global gitignore or merged into this project gitignore.  For a PyCharm
@@ -235,3 +232,4 @@ configs/development.yaml
 # Docker
 .dockerignore
 Dockerfile.dev
+discussion

README.md

@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://tinyurl.com/vllm-omni-meeting"><b>Weekly Meeting</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://tinyurl.com/vllm-omni-meeting"><b>Weekly Meeting</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
 </p>
 
 ---
@@ -30,7 +30,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 <p align="center">
   <picture>
-    <img alt="vllm-omni" src="./docs/source/architecture/omni-modality model architecture.png" width=55%>
+    <img alt="vllm-omni" src="./docs/source/architecture/omni-modality-model-architecture.png" width=55%>
   </picture>
 </p>
 
@@ -57,8 +57,6 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 
 Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn more.
 
-Visit our documentation to learn more.
-
 - [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
 - [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
 - [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)

docs/.nav.yml

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

docs/api/README.md

@@ -4,8 +4,8 @@
 
 Main entry points for vLLM-Omni inference and serving.
 
-- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
+- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
-- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniStageLLM][]
+- [vllm_omni.entrypoints.async_omni.AsyncOmniStageLLM][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
 - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
@@ -34,6 +34,8 @@ Engine classes for offline and online inference.
 - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
+- [vllm_omni.engine.OmniEngineCoreOutput][]
+- [vllm_omni.engine.OmniEngineCoreOutputs][]
 - [vllm_omni.engine.OmniEngineCoreRequest][]
 - [vllm_omni.engine.PromptEmbedsPayload][]
 - [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
@@ -48,6 +50,7 @@ Core scheduling and caching components.
 
 - [vllm_omni.core.dit_cache_manager.DiTCacheManager][]
 - [vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler][]
+- [vllm_omni.core.sched.generation_scheduler.GenerationScheduler][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
 - [vllm_omni.core.sched.scheduler.OmniScheduler][]
 
@@ -55,7 +58,7 @@ Core scheduling and caching components.
 
 Model execution components.
 
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.OmniOutput][]
+- [vllm_omni.model_executor.models.output_templates.OmniOutput][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniAudioFeatureInputs][]
@@ -72,6 +75,22 @@ Model execution components.
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2EmbeddingModel][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2ForCausalLM][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2Model][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_moe.Qwen3MoeForCausalLM][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni.Qwen3OmniMoeForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav.Qwen3OmniMoeCode2Wav][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniCodePredictorBaseModel][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniMoeTalkerCodePredictor][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeModel][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMForCausalLM][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMModel][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeConditionalGenerationMixin][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerMultiModalProcessor][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerProcessingInfo][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3Omni_VisionTransformer][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3_VisionPatchEmbed][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3_VisionPatchMerger][]
 
 ## Configuration
 
@@ -83,10 +102,15 @@ Configuration classes.
 
 Worker classes and model runners for distributed inference.
 
-- [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][]
-- [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
 - [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
 - [vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
+- [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
+- [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
 - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
+- [vllm_omni.worker.npu.npu_ar_model_runner.NPUARModelRunner][]
+- [vllm_omni.worker.npu.npu_ar_worker.NPUARWorker][]
+- [vllm_omni.worker.npu.npu_diffusion_model_runner.NPUDiffusionModelRunner][]
+- [vllm_omni.worker.npu.npu_diffusion_worker.NPUDiffusionWorker][]
+- [vllm_omni.worker.npu.npu_model_runner.OmniNPUModelRunner][]

docs/configuration/stage_configs.md

@@ -44,7 +44,7 @@ stage_args:
       trust_remote_code: true # Needed by huggingface config parsing
       engine_output_type: latent  # It claims that the stage will input latent hiddenstates besides token ids
       enable_prefix_caching: false # For request with hiddenstates output, the prefix caching is not supported now
-    is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmniLLM will use its tokenizer as default
+    is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmni will use its tokenizer as default
     final_output: true # If the stage has output as part of final outputs. If it is false, which means that the stage only works as a intermediate role.
     final_output_type: text # What is the final output type. It can be text and audio now.
     default_sampling_params: # sampling parameters for the stage. Their meaning aligns with vLLM.
@@ -206,7 +206,7 @@ Default: `false`
 
 ### `is_comprehension`
 
-Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmniLLM` will use its tokenizer as the default tokenizer for the entire pipeline.
+Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmni` will use its tokenizer as the default tokenizer for the entire pipeline.
 
 Default: `true`
 

docs/contributing/DOCS_GUIDE.md

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

docs/design/api_design_doc.md

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

docs/design/api_design_template.md

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

docs/design/vllm_omni_design.md

@@ -14,7 +14,7 @@ vLLM-Omni is a multi-modality extension for vLLM that supports non-autoregressiv
 
 ## Key Data Flow
 
-API Server --> OmniLLM/AsyncOmniLLM (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
+API Server --> OmniLLM/AsyncOmni (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)
 
@@ -40,7 +40,7 @@ graph TD
     B --> C{Detect --omni flag}
     C -->|Yes| D[Parse OmniConfig]
     C -->|No| E[Forward to vLLM CLI]
-    D --> F[Initialize AsyncOmniLLM]
+    D --> F[Initialize AsyncOmni]
     F --> G[Start omni Server]
     G --> H[Multi-stage Processing]
     E --> I[Standard vLLM Pipeline]
@@ -209,7 +209,7 @@ Similar to OmniLLM in offline inference, add some asynchronous processing, refer
 from vllm.v1.engine.async_llm import AsyncLLM
 
 
-class AsyncOmniLLM(AsyncLLM):
+class AsyncOmni(AsyncLLM):
     """Extended AsyncLLM supporting multiple engines and stage-based processing"""
 
     def __init__(self, stage_configs: List[StageConfig]):
@@ -449,15 +449,15 @@ AsyncLLM.add_request()
 ↓
 OutputProcessor.add_request()
 ↓
-RequestState.from_new_request() → 创建请求状态
+RequestState.from_new_request() → create request state
 
-AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() →  创建一个Background loop 持续从EngineCore获取输出
+AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() →  create a background loop that continuously pulls output from EngineCore
 ↓
-OutputProcessor.process_outputs() → 更新状态并处理输出
+OutputProcessor.process_outputs() → update state and process outputs
 ↓
-RequestState.make_request_output() → 转换为最终输出,格式为RequestOutput或者PoolingRequestOutput
+RequestState.make_request_output() → convert to the final output, formatted as RequestOutput or PoolingRequestOutput
 ↓
-RequestOutputCollector.put() → 推送到队列（AsyncLLM）
+RequestOutputCollector.put() → push to the queue (AsyncLLM)
 ```
 
 Need to add implementation for an existing method

docs/getting_started/installation/README.md

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

docs/getting_started/installation/npu.md

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

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

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

docs/mkdocs/hooks/generate_examples.py

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

docs/models/supported_models.md

@@ -17,6 +17,7 @@ th {
 
 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
 
@@ -33,3 +34,4 @@ th {
 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
+| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -1,55 +1,40 @@
-# Offline Inference Example of vLLM-Omni for Qwen2.5-Omni
+# Offline Example of vLLM-Omni for Qwen2.5-omni
 
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni>.
+Source <https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen2_5_omni>.
 
 
 ## 🛠️ Installation
 
-Please refer to [installation](../../../getting_started/installation/README.md).
+Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/README.md)
-
-## Run Offline inference with Qwen2.5-Omni.
-First, navigate to the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Inside the directory, `end2end.py` is a comprehensive demo suite for initializing a model instance of `Qwen/Qwen2.5-Omni-7B` and use it for various offline inference tasks.
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen2_5_omni/end2end.py"
-    ``````
-
-Below we also provide simple bash scripts to execute this file.
-### Single Prompt
-```bash
-bash run_single_prompt.sh
-```
-??? abstract "run_single_prompt.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
-    ``````
 
+## Run examples (Qwen2.5-omni)
 ### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit) and extract the prompts with `extract_prompts.py`
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
-??? abstract "extract_prompts.py"
-    ``````py
-    --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
-    ``````
 ```bash
 tar -xf <Your Download Path>/seedtts_testset.tar
 cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
 python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
   --input examples/offline_inference/qwen2_5_omni/meta.lst \
   --output examples/offline_inference/qwen2_5_omni/top100.txt \
-  --topk 10
+  --topk 100
+```
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
 ```
 Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
-??? abstract "run_multiple_prompts.sh"
+### Single Prompt
-    ``````sh
+Get into the example folder
-    --8<-- "examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
+```bash
-    ``````
+cd examples/offline_inference/qwen2_5_omni
+```
+Then run the command below.
+```bash
+bash run_single_prompt.sh
+```
 
 ### FAQ
 
@@ -58,3 +43,22 @@ If you encounter error about backend of librosa, try to install ffmpeg with comm
 sudo apt update
 sudo apt install ffmpeg
 ```
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/end2end.py"
+    ``````
+??? abstract "extract_prompts.py"
+    ``````py
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/extract_prompts.py"
+    ``````
+??? abstract "run_multiple_prompts.sh"
+    ``````sh
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
+    ``````
+??? abstract "run_single_prompt.sh"
+    ``````sh
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
+    ``````

docs/user_guide/examples/offline_inference/qwen3_omni.md

@@ -0,0 +1,64 @@
+# Offline Example of vLLM-Omni for Qwen3-omni
+
+Source <https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen3_omni>.
+
+
+## 🛠️ Installation
+
+Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/README.md)
+
+## Run examples (Qwen3-omni)
+### Multiple Prompts
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-omni README.md](https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen2_5_omni/README.md)
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen3_omni
+```
+Then run the command below.
+```bash
+bash run_multiple_prompts.sh
+```
+### Single Prompt
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen3_omni
+```
+Then run the command below.
+```bash
+bash run_single_prompt.sh
+```
+If you have not enough memory, you can set thinker with tensor parallel. Just run the command below.
+```bash
+bash run_single_prompt_tp.sh
+```
+
+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/end2end.py"
+    ``````
+??? abstract "qwen3_omni_moe_tp.yaml"
+    ``````yaml
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/qwen3_omni_moe_tp.yaml"
+    ``````
+??? abstract "run_multiple_prompts.sh"
+    ``````sh
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
+    ``````
+??? abstract "run_single_prompt.sh"
+    ``````sh
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_single_prompt.sh"
+    ``````
+??? abstract "run_single_prompt_tp.sh"
+    ``````sh
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
+    ``````

docs/user_guide/examples/offline_inference/qwen_image.md

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

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -1,43 +1,38 @@
-# Online Serving Example of vLLM-Omni for Qwen2.5-Omni
+# Online serving Example of vLLM-Omni for Qwen2.5-omni
 
-Source <https://github.com/vllm-project/vllm-omni/blob/main/examples/online_serving/README.md>.
+Source <https://github.com/vllm-project/vllm/tree/main/examples/online_serving/qwen2_5_omni>.
 
 
 ## 🛠️ Installation
 
-Please refer to [installation](../../../getting_started/installation/README.md).
+Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/examples/README.md)
 
-## Deploy Qwen/Qwen2.5-Omni-7B
+## Run examples (Qwen2.5-omni)
 
-First launch the OpenAI-compatible inference server
+Launch the server
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
+
 If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
-## Query the model
+
-Navigate to the example folder
+Get into the example folder
 ```bash
 cd examples/online_serving
 ```
-Query the model server with OpenAI Python API client:
+
+Send request via python
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```
-??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
+
-    ``````py
+Send request via curl
-    --8<-- "examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py"
-    ``````
-You can also query the model with `curl` command:
 ```bash
 bash run_curl_multimodal_generation.sh mixed_modalities
 ```
-??? abstract "run_curl_multimodal_generation.sh"
-    ``````py
-    --8<-- "examples/online_serving/run_curl_multimodal_generation.sh"
-    ``````
 
 ### FAQ
 
@@ -49,16 +44,17 @@ sudo apt install ffmpeg
 
 ## Run Local Web UI Demo
 
-You can also deploy a Gradio Web UI that allows users to interact with the model through a web browser. Below is an example on how to do so with `Qwen/Qwen2.5-Omni-7B`.
+This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo
 
-Install gradio with `uv pip install "gradio>=5.49.1,<6.0.0"`, then you can launch the web service built on AsyncOmni by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861
 ```
-Now you can interact with model via the web UI at `http://localhost:7861/` on your local browser.
+
+Then open `http://localhost:7861/` on your local browser to interact with the web UI.
 
 
 ### Options
@@ -78,3 +74,18 @@ python gradio_demo.py \
 - `--port`: Port for the Gradio server (default `7861`).
 - `--stage-configs-path`: Optional path to custom stage configs YAML.
 - `--share`: Set to expose a temporary public link via Gradio.
+
+## Example materials
+
+??? abstract "gradio_demo.py"
+    ``````py
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/gradio_demo.py"
+    ``````
+??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
+    ``````py
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    ``````
+??? abstract "run_curl_multimodal_generation.sh"
+    ``````sh
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh"
+    ``````

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -0,0 +1,95 @@
+# Online serving Example of vLLM-Omni for Qwen3-omni
+
+Source <https://github.com/vllm-project/vllm/tree/main/examples/online_serving/qwen3_omni>.
+
+
+## 🛠️ Installation
+
+Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/examples/README.md)
+
+## Run examples (Qwen3-Omni)
+
+Launch the server
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
+If you have custom stage configs file, launch the server with command below
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+```
+
+Get into the example folder
+```bash
+cd examples/online_serving
+```
+
+Send request via python
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
+```
+
+Send request via curl
+```bash
+bash run_curl_multimodal_generation.sh use_image
+```
+
+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
+
+## Run Local Web UI Demo
+
+This Web UI demo allows users to interact with the model through a web browser.
+
+### Running Gradio Demo
+
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
+
+```bash
+python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
+```
+
+Then open `http://localhost:7861/` on your local browser to interact with the web UI.
+
+
+### Options
+
+The gradio demo also supports running with an existing API server and can be customized with the following arguments.
+
+
+```bash
+python gradio_demo.py \
+    --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+    --use-api-server \
+    --api-base http://localhost:8091/v1 \
+    --ip 127.0.0.1 \
+    --port 7861
+```
+
+- `--model`: Model name
+- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmni locally.
+- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
+- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
+- `--port`: Port for Gradio server (default: 7861)
+- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
+- `--share`: Share the Gradio demo publicly (creates a public link)
+
+## Example materials
+
+??? abstract "gradio_demo.py"
+    ``````py
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/gradio_demo.py"
+    ``````
+??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
+    ``````py
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    ``````
+??? abstract "run_curl_multimodal_generation.sh"
+    ``````sh
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh"
+    ``````

examples/offline_inference/qwen2_5_omni/README.md

@@ -1,10 +1,10 @@
-# Offline Example of vLLM-Omni for Qwen2.5-omni
+# Offline Example of vLLM-Omni for Qwen2.5-Omni
 
 ## 🛠️ Installation
 
 Please refer to [README.md](../../../README.md)
 
-## Run examples (Qwen2.5-omni)
+## Run examples (Qwen2.5-Omni)
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
 ```bash
@@ -23,7 +23,7 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
-### Single Prompts
+### Single Prompt
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen2_5_omni

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -16,7 +16,7 @@ from vllm.multimodal.image import convert_image_mode
 from vllm.sampling_params import SamplingParams
 from vllm.utils import FlexibleArgumentParser
 
-from vllm_omni.entrypoints.omni_llm import OmniLLM
+from vllm_omni.entrypoints.omni import Omni
 
 SEED = 42
 
@@ -140,7 +140,7 @@ def main(args):
     model_name = "Qwen/Qwen2.5-Omni-7B"
     query_result = query_map[args.query_type]()
 
-    omni_llm = OmniLLM(
+    omni_llm = Omni(
         model=model_name,
         log_stats=args.enable_stats,
         log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
@@ -283,6 +283,19 @@ def parse_args():
         default=None,
         help="Path to a .txt file with one prompt per line (preferred).",
     )
+    parser.add_argument(
+        "--worker-backend",
+        type=str,
+        default="process",
+        choices=["process","ray"],
+        help="backend"
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+
 
     return parser.parse_args()
 

examples/offline_inference/qwen2_5_omni/extract_prompts.py

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

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

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

examples/offline_inference/qwen3_omni/README.md

@@ -1,12 +1,12 @@
-# Offline Example of vLLM-Omni for Qwen2.5-omni
+# Offline Example of vLLM-Omni for Qwen3-Omni
 
 ## 🛠️ Installation
 
 Please refer to [README.md](../../../README.md)
 
-## Run examples (Qwen3-omni)
+## Run examples (Qwen3-Omni)
 ### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-omni README.md](../qwen2_5_omni/README.md)
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-Omni README.md](../qwen2_5_omni/README.md)
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -15,7 +15,7 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
-### Single Prompts
+### Single Prompt
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -24,6 +24,10 @@ Then run the command below.
 ```bash
 bash run_single_prompt.sh
 ```
+If you have not enough memory, you can set thinker with tensor parallel. Just run the command below.
+```bash
+bash run_single_prompt_tp.sh
+```
 
 ### FAQ
 

examples/offline_inference/qwen3_omni/end2end.py

@@ -16,7 +16,7 @@ from vllm.assets.video import VideoAsset
 from vllm.multimodal.image import convert_image_mode
 from vllm.utils import FlexibleArgumentParser
 
-from vllm_omni.entrypoints.omni_llm import OmniLLM
+from vllm_omni.entrypoints.omni import Omni
 
 SEED = 42
 
@@ -127,8 +127,9 @@ def main(args):
     model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
     query_result = query_map[args.query_type]()
 
-    omni_llm = OmniLLM(
+    omni_llm = Omni(
         model=model_name,
+        stage_configs_path=args.stage_configs_path,
     )
 
     thinker_sampling_params = SamplingParams(
@@ -276,6 +277,12 @@ def parse_args():
         default=None,
         help="Path to a .txt file with one prompt per line (preferred).",
     )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to a stage configs file.",
+    )
 
     return parser.parse_args()
 
@@ -283,7 +290,3 @@ def parse_args():
 if __name__ == "__main__":
     args = parse_args()
     main(args)
-
-    # use examples:
-    # python end2end.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type text
-    # python end2end.py --model /custom_path/Qwen3-Omni-30B-A3B-Instruct --query-type use_video_only

examples/offline_inference/qwen3_omni/qwen3_omni_moe_tp.yaml

@@ -0,0 +1,94 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    runtime:
+       devices: "1"
+       max_batch_size: 1
+    engine_args:
+       model_stage: talker
+       model_arch: Qwen3OmniMoeForConditionalGeneration
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+       gpu_memory_utilization: 0.5
+       enforce_eager: true
+       trust_remote_code: true
+       engine_output_type: latent  # Output codec codes for code2wav
+      #  tensor_parallel_size: 2
+       enable_prefix_caching: false
+       distributed_executor_backend: "mp"
+       hf_config_name: talker_config
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+    # final_output: true
+    # final_output_type: text
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 2
+    runtime:
+      devices: "1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio  # Final output: audio waveform
+      gpu_memory_utilization: 0.1
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 1000000
+      hf_config_name: thinker_config
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1

examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh

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

examples/offline_inference/qwen_image/gradio_demo.py

@@ -5,6 +5,7 @@ import gradio as gr
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.platform_utils import detect_device_type
 
 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
     "1:1": (1328, 1328),
@@ -62,7 +63,7 @@ def get_omni(model_name: str) -> Omni:
 
 
 def build_demo(args: argparse.Namespace) -> gr.Blocks:
-    device = "cuda"
+    device = detect_device_type()
     omni = get_omni(args.model)
 
     def run_inference(

examples/offline_inference/qwen_image/text_to_image.py

@@ -7,6 +7,7 @@ from pathlib import Path
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.platform_utils import detect_device_type
 
 
 def parse_args() -> argparse.Namespace:
@@ -45,7 +46,7 @@ def parse_args() -> argparse.Namespace:
 
 def main():
     args = parse_args()
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = detect_device_type()
     generator = torch.Generator(device=device).manual_seed(args.seed)
 
     omni = Omni(model=args.model)

examples/offline_inference/tmp_test/README.md

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

examples/offline_inference/tmp_test/end2end.py

@@ -0,0 +1,264 @@
+import argparse
+import os
+import os as _os_env_toggle
+import random
+
+import numpy as np
+import soundfile as sf
+import torch
+from utils import make_omni_prompt
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+_os_env_toggle.environ["VLLM_USE_V1"] = "1"
+
+SEED = 42
+# Set all random seeds
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+
+# Make PyTorch deterministic
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+# Set environment variables for deterministic behavior
+os.environ["PYTHONHASHSEED"] = str(SEED)
+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", 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="[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).",
+    )
+    parser.add_argument(
+        "--worker-backend",
+        type=str,
+        default="process",
+        choices=["process","ray"],
+        help="backend"
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred)."
+        )
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    model_name = args.model
+    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, 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,
+        worker_backend=args.worker_backend,
+        ray_address=args.ray_address
+    )
+    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,
+    )
+    talker_sampling_params = SamplingParams(
+        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,
+    )
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        talker_sampling_params,
+        code2wav_sampling_params,
+    ]
+    import time
+    for i in range(1):
+        t1 = time.time()
+        prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
+        print(f"prompt:{prompt}")
+        omni_outputs = omni_llm.generate(prompt, sampling_params_list)
+        t2 = time.time()
+        print(f"==========> time:{t2-t1}")
+
+    # Determine output directory: prefer --output-dir; fallback to --output-wav
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    os.makedirs(output_dir, exist_ok=True)
+    for stage_outputs in omni_outputs:
+        if stage_outputs.final_output_type == "text":
+            for output in stage_outputs.request_output:
+                request_id = int(output.request_id)
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = 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 = int(output.request_id)
+                audio_tensor = output.multimodal_output["audio"]
+                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()

examples/offline_inference/tmp_test/extract_prompts.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+import argparse
+from typing import Optional
+
+
+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, 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/tmp_test/processing_omni.py

@@ -0,0 +1,367 @@
+from __future__ import annotations
+
+import base64
+import logging
+import math
+import os
+import time
+import warnings
+from functools import lru_cache
+from io import BytesIO
+
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+from torchvision import io, transforms
+from torchvision.transforms import InterpolationMode
+
+logger = logging.getLogger(__name__)
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+VIDEO_MIN_PIXELS = 128 * 28 * 28
+VIDEO_MAX_PIXELS = 768 * 28 * 28
+VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
+FRAME_FACTOR = 2
+FPS = 2.0
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+
+temporal_patch_size = 2
+spatial_patch_size = 14
+spatial_merge_size = 2
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """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'."""
+    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]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    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)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        image_obj = Image.open(requests.get(image, stream=True).raw)
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        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}")
+    image = image_obj.convert("RGB")
+    # resize
+    if "resized_height" in ele and "resized_width" in ele:
+        resized_height, resized_width = smart_resize(
+            ele["resized_height"],
+            ele["resized_width"],
+            factor=size_factor,
+        )
+    else:
+        width, height = image.size
+        min_pixels = ele.get("min_pixels", MIN_PIXELS)
+        max_pixels = ele.get("max_pixels", MAX_PIXELS)
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=size_factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+    image = image.resize((resized_width, resized_height))
+
+    return image
+
+
+def smart_nframes(
+    ele: dict,
+    total_frames: int,
+    video_fps: int | float,
+) -> int:
+    """calculate the number of frames for video used for model inputs.
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+            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.
+        total_frames (int): the original total number of frames of the video.
+        video_fps (int | float): the original fps of the video.
+
+    Raises:
+        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
+
+    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`"
+    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)
+        max_frames = floor_by_factor(
+            ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
+            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}.")
+    return nframes
+
+
+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_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).
+    """
+    video_path = ele["video"]
+    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.")
+        if "file://" in video_path:
+            video_path = video_path[7:]
+    st = time.time()
+    video, audio, info = io.read_video(
+        video_path,
+        start_pts=ele.get("video_start", 0.0),
+        end_pts=ele.get("video_end", None),
+        pts_unit="sec",
+        output_format="TCHW",
+    )
+    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"
+    )
+    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
+
+
+def is_decord_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("decord") is not None
+
+
+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_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:
+        raise NotImplementedError("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")
+    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
+    return video, total_duration, nframes
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if FORCE_QWENVL_VIDEO_READER is not None:
+        video_reader_backend = FORCE_QWENVL_VIDEO_READER
+    elif is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    # 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]:
+    if isinstance(ele["video"], str):
+        video_reader_backend = get_video_reader_backend()
+        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
+        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),
+        )
+        max_pixels = ele.get("max_pixels", max_pixels)
+        # min_pixels = (factor ** 2) * 52
+        # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
+        if "resized_height" in ele and "resized_width" in ele:
+            resized_height, resized_width = smart_resize(
+                ele["resized_height"],
+                ele["resized_width"],
+                factor=image_factor,
+            )
+        else:
+            resized_height, resized_width = smart_resize(
+                height,
+                width,
+                factor=image_factor,
+                min_pixels=min_pixels,
+                max_pixels=max_pixels,
+            )
+        video = transforms.functional.resize(
+            video,
+            [resized_height, resized_width],
+            interpolation=InterpolationMode.BICUBIC,
+            antialias=True,
+        ).float()
+        return video, total_dur, nframes, second_per_grid
+    else:
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        images = [
+            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)
+        if len(images) < nframes:
+            images.extend([images[-1]] * (nframes - len(images)))
+        return images, None, None, None
+
+
+def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations = [conversations]
+    for conversation in conversations:
+        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")
+                    ):
+                        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]:
+    vision_infos = extract_vision_info(conversations)
+    # Read images or videos
+    image_inputs = []
+    video_inputs = []
+    for vision_info in vision_infos:
+        if "image" in vision_info or "image_url" in vision_info:
+            image_inputs.append(fetch_image(vision_info))
+        elif "video" in vision_info:
+            video_inputs.append(fetch_video(vision_info))
+        else:
+            raise ValueError("image, image_url or video should in content.")
+    if len(image_inputs) == 0:
+        image_inputs = None
+    if len(video_inputs) == 0:
+        video_inputs = None
+    return image_inputs, video_inputs

examples/offline_inference/tmp_test/ray_head.sh

@@ -0,0 +1 @@
+ray start --head --port=6399 --num-gpus=1

examples/offline_inference/tmp_test/ray_worker.sh

@@ -0,0 +1 @@
+ray start --address=10.90.67.86:6399

examples/offline_inference/tmp_test/run_multiple_prompts.sh

@@ -0,0 +1,8 @@
+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/tmp_test/run_single_prompt.sh

@@ -0,0 +1,11 @@
+export PYTHONPATH=/workspace/omni/vllm-omni/:$PYTHONPATH
+CUDA_VISIBLE_DEVICES=4 python end2end.py --model /workspace/Qwen2.5-Omni-7B/ \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+				 --worker-backend ray \
+				 --ray-address auto \
+                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."

examples/offline_inference/tmp_test/utils.py

@@ -0,0 +1,312 @@
+import tempfile
+from typing import Optional, Union
+from urllib.request import urlopen
+
+import librosa
+import requests
+import resampy
+import soundfile as sf
+import torch
+import torchvision.io
+from processing_omni import fetch_image, fetch_video
+from transformers import AutoConfig, AutoProcessor
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+# 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."
+                ),
+            }
+        ],
+    }
+
+
+def resample_wav_to_16khz(input_filepath):
+    data, original_sample_rate = sf.read(input_filepath)
+    # Only use the first channel
+    if len(data.shape) > 1:
+        data = data[:, 0]
+    # resample to 16kHz
+    data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000)
+    return data_resampled
+
+
+def fetch_and_read_video(args, video_url: str, fps=2):
+    def read_video_with_torchvision(video_file_name: str):
+        video, audio, info = torchvision.io.read_video(
+            video_file_name,
+            start_pts=0.0,
+            end_pts=None,
+            pts_unit="sec",
+            output_format="TCHW",
+        )
+
+        total_frames, video_fps = video.size(0), info["video_fps"]
+        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
+        grid_timestamps = frame_timestamps[::2]
+        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
+
+        idx = torch.linspace(0, video.size(0) - 1, nframes).round().long()
+        video = video[idx]
+
+        if args.legacy_omni_video:
+            return [video, total_duration, nframes, second_per_grid.item()]
+        else:
+            return video
+
+    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})
+        if total_duration is None and nframes is None:
+            nframes = len(video)
+            total_duration = 0.5 * nframes
+            second_per_grid = 1.0
+        if args.legacy_omni_video:
+            return [video, total_duration, nframes, second_per_grid]
+        else:
+            return video
+
+    def read_video(video_file_name: str):
+        if args.use_torchvision:
+            return read_video_with_torchvision(video_file_name)
+        else:
+            return read_video_with_transformers(video_file_name)
+
+    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}"
+            )
+
+            temp_video_file.write(urlopen(video_url).read())
+            temp_video_file_path = temp_video_file.name
+            video_file_name = temp_video_file_path
+            return read_video(video_file_name)
+    else:
+        video_file_name = video_url
+        return read_video(video_file_name)
+
+
+def make_inputs_qwen2_omni(
+    args,
+    messages: list[dict[str, Union[str, list[dict[str, str]]]]],
+    use_audio_in_video: Optional[bool] = False,
+    tokenize: bool = False,
+) -> Union[OmniTokensPrompt, TextPrompt]:
+    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 = "Qwen2_5OmniModel" not in config.architectures
+        else:
+            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"])
+        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"
+                else:
+                    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:
+                        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(ele[audio_key]))
+                else:
+                    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:
+                        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(ele[audio_key], sr=16000)[0])
+                    videos.append(fetch_and_read_video(args, audio_key))
+                else:
+                    raise ValueError(f"Unknown ele {ele}")
+                # insert a audio after the video
+                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())
+                        temp_video_file_path = temp_video_file.name
+                        videos.append(fetch_and_read_video(args, temp_video_file))
+                        ele["video"] = temp_video_file_path
+                else:
+                    video_key = "video"
+                    videos.append(fetch_and_read_video(args, ele[video_key]))
+            elif ele["type"] == "image" or ele["type"] == "image_url":
+                images.append(fetch_image(ele))
+
+            # move to the next content
+            index += 1
+
+    prompt = processor.apply_chat_template(
+        messages,
+        tokenize=tokenize,
+        add_generation_prompt=True,
+        add_vision_id=True,
+    )
+
+    audios = audios if len(audios) > 0 else None
+    images = images if len(images) > 0 else None
+    videos = videos if len(videos) > 0 else None
+
+    multi_modal_data = {}
+    if audios:
+        multi_modal_data["audio"] = audios
+    if images:
+        multi_modal_data["image"] = images
+    if videos:
+        multi_modal_data["video"] = videos
+
+    if isinstance(prompt, list) and isinstance(prompt[0], (list, str)):
+        prompt = prompt[0]
+
+    if tokenize:
+        return OmniTokensPrompt(
+            prompt_token_ids=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+    else:
+        return TextPrompt(
+            prompt=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+
+
+def make_text_prompt(args, prompt):
+    messages = [
+        get_system_prompt(),
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt},
+            ],
+        },
+    ]
+
+    prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize)
+    return 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",
+            "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(
+        args,
+        messages,
+        use_audio_in_video=True,
+        tokenize=args.tokenize,
+    )
+    return prompt
+
+
+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":
+        prompt = make_audio_in_video_v2_prompt(args)
+    else:
+        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
+    return prompt

examples/online_serving/qwen2_5_omni/README.md

@@ -1,10 +1,10 @@
-# Online serving Example of vLLM-Omni for Qwen2.5-omni
+# Online serving Example of vLLM-Omni for Qwen2.5-Omni
 
 ## 🛠️ Installation
 
-Please refer to [README.md](../../README.md)
+Please refer to [README.md](../../../README.md)
 
-## Run examples (Qwen2.5-omni)
+## Run examples (Qwen2.5-Omni)
 
 Launch the server
 ```bash
@@ -18,7 +18,7 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to
 
 Get into the example folder
 ```bash
-cd examples/online_serving
+cd examples/online_serving/qwen2_5_omni
 ```
 
 Send request via python
@@ -45,7 +45,7 @@ This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo
 
-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861

examples/online_serving/qwen2_5_omni/gradio_demo.py

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

examples/online_serving/qwen3_omni/README.md

@@ -1,8 +1,8 @@
-# Online serving Example of vLLM-omni for Qwen2.5-omni
+# Online serving Example of vLLM-Omni for Qwen3-omni
 
 ## 🛠️ Installation
 
-Please refer to [README.md](../../README.md)
+Please refer to [README.md](../../../README.md)
 
 ## Run examples (Qwen3-Omni)
 
@@ -18,17 +18,17 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-p
 
 Get into the example folder
 ```bash
-cd examples/online_serving
+cd examples/online_serving/qwen3_omni
 ```
 
 Send request via python
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
+python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
 ```
 
 Send request via curl
 ```bash
-bash run_curl_multimodal_generation.sh mixed_modalities
+bash run_curl_multimodal_generation.sh use_image
 ```
 
 ### FAQ
@@ -45,7 +45,7 @@ This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo
 
-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
@@ -69,7 +69,7 @@ python gradio_demo.py \
 ```
 
 - `--model`: Model name
-- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmniLLM locally.
+- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmni locally.
 - `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
 - `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
 - `--port`: Port for Gradio server (default: 7861)

examples/online_serving/qwen3_omni/gradio_demo.py

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

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -197,7 +197,6 @@ def run_multimodal_generation(args) -> None:
 
     count = 0
     for choice in chat_completion.choices:
-        print(f"choice: {choice}")
         if choice.message.audio:
             audio_data = base64.b64decode(choice.message.audio.data)
             audio_file_path = f"audio_{count}.wav"

examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh

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

mkdocs.yml

@@ -58,6 +58,7 @@ theme:
 hooks:
   - docs/mkdocs/hooks/generate_api_readme.py
   - docs/mkdocs/hooks/url_schemes.py
+  - docs/mkdocs/hooks/generate_examples.py
 
 # Plugins
 plugins:
@@ -87,6 +88,7 @@ plugins:
       show_full_namespace: false  # Show only module name, not full path
       exclude:
         - "re:vllm_omni\\._.*"  # Internal modules
+        - "vllm_omni.diffusion.models.qwen_image"  # avoid importing vllm in mkdocs building
   - mkdocstrings:
       handlers:
         python:

tests/__init__.py

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

tests/test_omni_llm.py

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

vllm_omni/__init__.py

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

vllm_omni/core/sched/diffusion_scheduler.py

@@ -4,9 +4,9 @@ from typing import Optional
 
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.request_queue import create_request_queue
-from vllm.v1.core.sched.scheduler import EngineCoreOutputs, Request, RequestStatus, SchedulerOutput, SpecDecodingStats
+from vllm.v1.core.sched.scheduler import Request, RequestStatus, SchedulerOutput, SpecDecodingStats
 from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
 
 from vllm_omni.core.sched.output import OmniNewRequestData
 from vllm_omni.core.sched.scheduler import OmniScheduler

vllm_omni/core/sched/generation_scheduler.py

@@ -5,14 +5,13 @@ 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.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
 
 from vllm_omni.core.sched.output import OmniNewRequestData
 from vllm_omni.core.sched.scheduler import OmniScheduler

vllm_omni/core/sched/scheduler.py

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

vllm_omni/diffusion/diffusion_engine.py

@@ -9,7 +9,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.registry import get_diffusion_post_process_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.scheduler import scheduler
-from vllm_omni.diffusion.worker.gpu_worker import WorkerProc
+from vllm_omni.utils.platform_utils import get_diffusion_worker_class
 
 logger = init_logger(__name__)
 
@@ -76,6 +76,9 @@ class DiffusionEngine:
         mp.set_start_method("spawn", force=True)
         processes = []
 
+        # Get the appropriate worker class for current device
+        worker_proc = get_diffusion_worker_class()
+
         # Launch all worker processes
         scheduler_pipe_readers = []
         scheduler_pipe_writers = []
@@ -84,7 +87,7 @@ class DiffusionEngine:
             reader, writer = mp.Pipe(duplex=False)
             scheduler_pipe_writers.append(writer)
             process = mp.Process(
-                target=WorkerProc.worker_main,
+                target=worker_proc.worker_main,
                 args=(
                     i,  # rank
                     od_config,

vllm_omni/diffusion/distributed/utils.py

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

vllm_omni/diffusion/models/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Diffusion model implementations."""

vllm_omni/diffusion/models/qwen_image/__init__.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Qwen Image diffusion model components."""
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import (
+    QwenImagePipeline,
+    get_qwen_image_post_process_func,
+)
+from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
+    QwenImageTransformer2DModel,
+)
+
+__all__ = [
+    "QwenImagePipeline",
+    "QwenImageTransformer2DModel",
+    "get_qwen_image_post_process_func",
+]

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -240,19 +240,32 @@ class QwenImagePipeline(
         prefix: str = "",
     ):
         super().__init__()
+        self.od_config = od_config
         self.device = get_local_device()
         model = od_config.model
-        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model, subfolder="scheduler")
+
+        # Check if model is a local path
+        local_files_only = os.path.exists(model)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
         logger.info("Loaded Qwen-Image scheduler successfully")
 
         with set_default_torch_dtype(torch.bfloat16):
-            self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(model, subfolder="text_encoder")
+            self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model, subfolder="text_encoder", local_files_only=local_files_only
+            )
             logger.info("Loaded Qwen-Image text encoder successfully")
-            self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae").to(self.device)
+            self.vae = AutoencoderKLQwenImage.from_pretrained(
+                model, subfolder="vae", local_files_only=local_files_only
+            ).to(self.device)
             logger.info("Loaded Qwen-Image VAE successfully")
             self.transformer = QwenImageTransformer2DModel()
             logger.info("Initialized Qwen-Image transformer successfully.")
-            self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer")
+            self.tokenizer = Qwen2Tokenizer.from_pretrained(
+                model, subfolder="tokenizer", local_files_only=local_files_only
+            )
             logger.info("Loaded Qwen-Image tokenizer successfully.")
 
         self.stage = None
@@ -392,8 +405,6 @@ class QwenImagePipeline(
         Args:
             prompt (`str` or `list[str]`, *optional*):
                 prompt to be encoded
-            device: (`torch.device`):
-                torch device
             num_images_per_prompt (`int`):
                 number of images that should be generated per prompt
             prompt_embeds (`torch.Tensor`, *optional*):
@@ -765,10 +776,14 @@ class QwenImagePipeline(
                     for name, tensor in state_dict.items():
                         yield name, tensor
 
-        model_name = "Qwen/Qwen-Image"
         try:
-            # Use the arguments as verified by the user
+            # Get model path from config or download from HF
-            model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+            model_name = self.od_config.model if hasattr(self, "od_config") else "Qwen/Qwen-Image"
+            if os.path.exists(model_name):
+                model_path = model_name
+            else:
+                model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+
             transformer_path = os.path.join(model_path, "transformer")
             self.transformer.load_weights(weight_iterator(transformer_path))
 

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -113,10 +113,12 @@ class QwenEmbedRope(nn.Module):
         # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
         self.scale_rope = scale_rope
 
-    def rope_params(self, index, dim, theta=10000):
+    def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000):
         """
         Args:
-            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
+            index (`torch.Tensor`): [0, 1, 2, 3] 1D Tensor representing the position index of the token
+            dim (`int`): Dimension for the rope parameters
+            theta (`int`): Theta parameter for rope
         """
         assert dim % 2 == 0
         freqs = torch.outer(

vllm_omni/diffusion/worker/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Worker classes for diffusion models."""
+
+from vllm_omni.diffusion.worker.gpu_worker import GPUWorker, WorkerProc
+
+__all__ = ["GPUWorker", "WorkerProc"]

vllm_omni/diffusion/worker/npu/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NPU worker classes for diffusion models."""
+
+from vllm_omni.diffusion.worker.npu.npu_worker import NPUWorker, NPUWorkerProc
+
+__all__ = ["NPUWorker", "NPUWorkerProc"]

vllm_omni/diffusion/worker/npu/npu_worker.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import multiprocessing as mp
+import os
+
+import torch
+import zmq
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
+from vllm.distributed.parallel_state import (
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.registry import initialize_model
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+logger = init_logger(__name__)
+
+
+class NPUWorker:
+    """
+    A worker that executes the model on a single NPU.
+    """
+
+    def __init__(
+        self,
+        local_rank: int,
+        rank: int,
+        od_config: OmniDiffusionConfig,
+    ):
+        self.local_rank = local_rank
+        self.rank = rank
+        self.od_config = od_config
+        self.pipeline = None
+
+        self.init_device_and_model()
+
+    def init_device_and_model(self) -> None:
+        """Initialize the device and load the model."""
+        world_size = self.od_config.num_gpus
+        rank = self.rank
+        # Set environment variables for distributed initialization
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(self.od_config.master_port)
+        os.environ["LOCAL_RANK"] = str(self.local_rank)
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+
+        device = torch.device(f"npu:{rank}")
+        torch.npu.set_device(device)
+
+        # hack
+        vllm_config = VllmConfig()
+        vllm_config.parallel_config.tensor_parallel_size = self.od_config.num_gpus
+        set_current_vllm_config(vllm_config)
+
+        init_distributed_environment(world_size=world_size, rank=rank)
+        initialize_model_parallel(tensor_model_parallel_size=world_size)
+
+        with device:
+            self.pipeline = initialize_model(self.od_config)
+            self.pipeline.load_weights()
+            self.pipeline.eval()
+        logger.info(f"Worker {self.rank}: Initialized device, model, and distributed environment.")
+        logger.info(f"Worker {self.rank}: Model loaded successfully.")
+
+    @torch.inference_mode()
+    def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
+        """
+        Execute a forward pass.
+        """
+        assert self.pipeline is not None
+        # TODO: dealing with first req for now
+        req = reqs[0]
+        output = self.pipeline.forward(req)
+        return output
+
+
+class NPUWorkerProc:
+    """Wrapper that runs one Worker in a separate process."""
+
+    def __init__(
+        self,
+        od_config: OmniDiffusionConfig,
+        gpu_id: int,
+        broadcast_handle,
+    ):
+        self.od_config = od_config
+
+        # Inter-process Communication
+        self.context = zmq.Context(io_threads=2)
+
+        # Initialize MessageQueue reader from handle
+        self.mq = MessageQueue.create_from_handle(broadcast_handle, gpu_id)
+
+        self.result_mq = None
+        self.result_mq_handle = None
+
+        # Setup result sender (only for rank 0 for now, or whoever needs to reply)
+        # Assuming only rank 0 replies to scheduler as per original logic
+        if gpu_id == 0:
+            # Create MessageQueue for results (1 writer -> 1 reader)
+            # We assume the reader (SyncScheduler) will act as rank 0
+            self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
+            self.result_mq_handle = self.result_mq.export_handle()
+            logger.info(f"Worker {gpu_id} created result MessageQueue")
+
+        assert od_config.master_port is not None
+        worker = NPUWorker(
+            local_rank=gpu_id,
+            rank=gpu_id,
+            od_config=od_config,
+        )
+        self.worker = worker
+        self.gpu_id = gpu_id
+        self._running = True
+
+    def return_result(self, output: DiffusionOutput):
+        """
+        replies to client, only on rank 0
+        """
+        if self.result_mq is not None:
+            self.result_mq.enqueue(output)
+
+    def recv_reqs(self):
+        """
+        Receive requests from broadcast queue
+        """
+        return self.mq.dequeue()
+
+    # TODO: queueing, cancellation
+    def worker_busy_loop(self) -> None:
+        """Main busy loop for Multiprocessing Workers"""
+
+        logger.info(f"Worker {self.gpu_id} ready to receive requests via shared memory")
+
+        while self._running:
+            reqs = None
+            # 1: receive requests
+            try:
+                reqs = self.recv_reqs()
+            except Exception as e:
+                logger.error(
+                    f"Error receiving requests in scheduler event loop: {e}",
+                    exc_info=True,
+                )
+                continue
+
+            # 2: execute, make sure a reply is always sent
+            try:
+                output = self.worker.execute_model(reqs, self.od_config)
+            except Exception as e:
+                logger.error(
+                    f"Error executing forward in event loop: {e}",
+                    exc_info=True,
+                )
+                output = DiffusionOutput(error=str(e))
+
+            try:
+                self.return_result(output)
+            except zmq.ZMQError as e:
+                # Reply failed; log and keep loop alive to accept future requests
+                logger.error(f"ZMQ error sending reply: {e}")
+                continue
+
+        logger.info("event loop terminated.")
+        # if self.result_sender is not None:
+        #     self.result_sender.close()
+        self.context.term()
+
+    @staticmethod
+    def worker_main(
+        rank: int,
+        od_config: OmniDiffusionConfig,
+        pipe_writer: mp.connection.Connection,
+        broadcast_handle,
+    ) -> None:
+        """Worker initialization and execution loops."""
+
+        worker_proc = NPUWorkerProc(
+            od_config,
+            gpu_id=rank,
+            broadcast_handle=broadcast_handle,
+        )
+        logger.info(f"Worker {rank}: Scheduler loop started.")
+        pipe_writer.send(
+            {
+                "status": "ready",
+                "result_handle": worker_proc.result_mq_handle if rank == 0 else None,
+            }
+        )
+        worker_proc.worker_busy_loop()
+        logger.info(f"Worker {rank}: Shutdown complete.")

vllm_omni/distributed/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .connectors import (
+    ConnectorSpec, OmniTransferConfig,
+    OmniConnectorBase, InMemoryOmniConnector,
+    OmniConnectorFactory, MooncakeConnector,
+    load_omni_transfer_config, create_simple_config, create_mooncake_config
+)
+
+__all__ = [
+    # Config
+    "ConnectorSpec",
+    "OmniTransferConfig",
+
+    # Connectors
+    "OmniConnectorBase",
+    "InMemoryOmniConnector",
+    "OmniConnectorFactory",
+    "MooncakeConnector",
+
+    # Utilities
+    "load_omni_transfer_config",
+    "create_simple_config",
+    "create_mooncake_config",
+]

vllm_omni/distributed/connectors/__init__.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .config import ConnectorSpec, OmniTransferConfig
+from .base import OmniConnectorBase, InMemoryOmniConnector
+from .factory import OmniConnectorFactory
+from .mooncake_connector import MooncakeConnector
+from .utils import (
+    load_omni_transfer_config,
+    initialize_connectors_from_config,
+    create_simple_config,
+    create_mooncake_config,
+    get_connectors_config_for_stage,
+    initialize_orchestrator_connectors,
+    get_stage_connector_config,
+    build_stage_connectors,
+)
+
+__all__ = [
+    # Config
+    "ConnectorSpec",
+    "OmniTransferConfig",
+
+    # Base classes and implementations
+    "OmniConnectorBase",
+    "InMemoryOmniConnector",
+
+    # Factory
+    "OmniConnectorFactory",
+
+    # Specific implementations
+    "MooncakeConnector",
+
+    # Utilities
+    "load_omni_transfer_config",
+    "initialize_connectors_from_config",
+    "create_simple_config",
+    "create_mooncake_config",
+    "get_connectors_config_for_stage",
+
+    # Manager helpers
+    "initialize_orchestrator_connectors",
+    "get_stage_connector_config",
+    "build_stage_connectors",
+]

vllm_omni/distributed/connectors/adapter.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# temporary for compatibility with vllm_omni.entrypoints.omni_stage.py 
+# and vllm_omni.entrypoints.omni_llm.py
+
+import time
+import logging
+from typing import Any, Optional, Dict, Callable
+
+logger = logging.getLogger(__name__)
+
+
+def try_send_via_connector(
+    connector: Any,
+    stage_id: int,
+    next_stage_id: int,
+    req_id: str,
+    next_inputs: Any,
+    sampling_params: Any,
+    original_prompt: Any,
+    next_stage_queue_submit_fn: Callable[[Dict[str, Any]], None],
+    metrics: Any
+) -> bool:
+    """
+    Attempts to send data via OmniConnector.
+    Returns True if successful, False otherwise.
+    Encapsulates the logic of preparing payload, sending via connector, 
+    sending notification, and recording metrics.
+    """
+    try:
+        t0 = time.time()
+
+        # Prepare data for connector
+        payload_data = {
+            "engine_inputs": next_inputs,
+            "sampling_params": sampling_params,
+            "metadata": {
+                "original_prompt": original_prompt,
+                "stage_transition": f"{stage_id}->{next_stage_id}",
+                "timestamp": time.time()
+            }
+        }
+
+        # Send data via connector
+        success, serialized_size = connector.put(str(stage_id), str(next_stage_id), str(req_id), payload_data)
+        
+        if success:
+            # Send lightweight notification via queue
+            notify_payload = {
+                "request_id": req_id,
+                "sampling_params": sampling_params,
+                "from_connector": True,
+                "from_stage": str(stage_id),
+                "to_stage": str(next_stage_id),
+                "sent_ts": time.time(),
+            }
+            next_stage_queue_submit_fn(notify_payload)
+
+            t1 = time.time()
+            tx_ms = (t1 - t0) * 1000.0
+            
+            metrics.on_forward(
+                stage_id,
+                next_stage_id,
+                req_id,
+                serialized_size,  # Use size from connector
+                float(tx_ms),
+                True,  # Mark as using connector
+            )
+            return True
+        else:
+            # If put returned False, we let the caller handle fallback
+            return False
+
+    except Exception as e:
+        logger.warning(
+            "[Orchestrator] OmniConnector failed for req %s: %s; falling back to queue",
+            req_id,
+            e,
+        )
+        return False
+
+
+def try_recv_via_connector(
+    task: Dict[str, Any], 
+    connectors: Dict[Any, Any], 
+    stage_id: int,
+) -> tuple[Any, Optional[Dict[str, Any]]]:
+    """
+    Attempts to resolve input data from either connector or IPC.
+    Returns (engine_inputs, rx_metrics) or (None, None) if failed/skipped.
+    """
+    rid = task["request_id"]
+    
+    if task.get("from_connector"):
+        from_stage = task.get("from_stage")
+        to_stage = str(stage_id)
+
+        if not from_stage:
+            logger.error(
+                "[Stage-%s] 'from_connector' is true but 'from_stage' is missing for request %s", stage_id, rid
+            )
+            return None, None
+
+        # Get connector for this edge
+        connector_key = (from_stage, to_stage)
+        connector = connectors.get(connector_key)
+
+        if connector:
+            try:
+                # Get data from connector with timeout
+                _t_start = time.time()
+                payload_data = connector.get(from_stage, to_stage, str(rid))
+                _t_end = time.time()
+                
+                if payload_data and isinstance(payload_data, dict):
+                    ein = payload_data.get("engine_inputs")
+                    # Use connector's serialization size for consistency
+                    serialized_size = len(connector.serialize_obj(payload_data))
+                    decode_ms = (_t_end - _t_start) * 1000.0
+                    
+                    rx_metrics = {
+                        "rx_decode_time_ms": decode_ms,
+                        "rx_transfer_bytes": serialized_size
+                    }
+                    return ein, rx_metrics
+                else:
+                    logger.error(
+                        "[Stage-%s] Failed to get data from connector for request %s or payload is empty", stage_id, rid
+                    )
+                    return None, None
+            except Exception as e:
+                logger.error(
+                    "[Stage-%s] Error retrieving data from connector for request %s: %s", stage_id, rid, e
+                )
+                return None, None
+        else:
+            logger.error(
+                "[Stage-%s] No connector found for edge %s -> %s for request %s", stage_id, from_stage, to_stage, rid
+            )
+            return None, None
+    else:
+        # Data comes from queue as usual
+        #return maybe_load_from_ipc_with_metrics_fn(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+        return None, None
+

vllm_omni/distributed/connectors/base.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class OmniConnectorBase(ABC):
+    """Base class for all OmniConnectors."""
+
+    @abstractmethod
+    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
+        """Store Python object, internal serialization handled by connector.
+
+        Args:
+            from_stage: Source stage identifier
+            to_stage: Destination stage identifier
+            request_id: Unique request identifier
+            data: Python object to store
+
+        Returns:
+            tuple: (success: bool, serialized_size: int)
+        """
+        pass
+
+    @abstractmethod
+    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
+        """Retrieve Python object, internal deserialization handled by connector.
+
+        Args:
+            from_stage: Source stage identifier
+            to_stage: Destination stage identifier
+            request_id: Unique request identifier
+
+        Returns:
+            Python object if found, None otherwise
+        """
+        pass
+
+    @abstractmethod
+    def cleanup(self, request_id: str) -> None:
+        """Clean up resources for a request."""
+        pass
+
+    @abstractmethod
+    def health(self) -> dict[str, Any]:
+        """Return health status and metrics."""
+        pass
+
+    @staticmethod
+    def serialize_obj(obj: Any) -> bytes:
+        """Serialize a Python object to bytes using torch.save (consistent with stage_utils)."""
+        import io
+        import torch
+        buf = io.BytesIO()
+        torch.save(obj, buf)
+        return buf.getvalue()
+
+    @staticmethod
+    def deserialize_obj(data: bytes) -> Any:
+        """Deserialize bytes to Python object using torch.load (consistent with stage_utils)."""
+        import io
+        import torch
+        return torch.load(io.BytesIO(data), map_location="cpu", weights_only=False)
+
+
+class InMemoryOmniConnector(OmniConnectorBase):
+    """Simple in-memory connector for testing and co-located stages."""
+
+    def __init__(self):
+        self._storage: dict[str, Any] = {}
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "errors": 0,
+        }
+
+    def _make_key(self, from_stage: str, to_stage: str, request_id: str) -> str:
+        return f"{from_stage}:{to_stage}:{request_id}"
+
+    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
+        try:
+            key = self._make_key(from_stage, to_stage, request_id)
+            self._storage[key] = data
+            serialized_size = len(self.serialize_obj(data))
+            self._metrics["puts"] += 1
+            logger.debug(f"InMemoryConnector: stored {key}")
+            return True, serialized_size
+        except Exception as e:
+            self._metrics["errors"] += 1
+            logger.error(f"InMemoryConnector put failed: {e}")
+            return False, 0
+
+    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
+        try:
+            key = self._make_key(from_stage, to_stage, request_id)
+            data = self._storage.get(key)
+            if data is not None:
+                self._metrics["gets"] += 1
+                return data
+            return None
+        except Exception as e:
+            logger.error(f"InMemoryConnector get failed: {e}")
+            return None
+
+    def cleanup(self, request_id: str) -> None:
+        keys_to_remove = [k for k in self._storage.keys() if k.endswith(f":{request_id}")]
+        for key in keys_to_remove:
+            del self._storage[key]
+        logger.debug(f"InMemoryConnector: cleaned up {len(keys_to_remove)} keys for {request_id}")
+
+    def health(self) -> dict[str, Any]:
+        return {
+            "status": "healthy",
+            "storage_size": len(self._storage),
+            **self._metrics,
+        }

vllm_omni/distributed/connectors/config.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ConnectorSpec:
+    """Specification for a connector instance."""
+    name: str  # e.g., "MooncakeConnector", "InMemoryConnector"
+    extra: dict[str, Any] = field(default_factory=dict)  # backend-specific config
+
+
+@dataclass
+class OmniTransferConfig:
+    """Top-level configuration for OmniConnector system."""
+    # Direct mapping: (from_stage, to_stage) -> connector
+    connectors: dict[tuple[str, str], ConnectorSpec] = field(default_factory=dict)
+    default_connector: Optional[ConnectorSpec] = None
+
+    def get_connector_for_edge(self, from_stage: str, to_stage: str) -> Optional[ConnectorSpec]:
+        """Get connector spec for a specific edge."""
+        edge_key = (from_stage, to_stage)
+        return self.connectors.get(edge_key, self.default_connector)
+
+    def has_connector_for_edge(self, from_stage: str, to_stage: str) -> bool:
+        """Check if there's a connector configured for the edge."""
+        return self.get_connector_for_edge(from_stage, to_stage) is not None

vllm_omni/distributed/connectors/connector_tests/run_single_prompt.sh

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

vllm_omni/distributed/connectors/connector_tests/test_connector_integration.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""
+Test OmniConnector integration with OmniLLM.
+"""
+
+import sys
+import os
+sys.path.insert(0, os.path.dirname(__file__))
+
+from vllm_omni.distributed.connectors import create_mooncake_config
+
+def test_connector_config():
+    """Test connector configuration creation."""
+    print("Testing connector configuration...")
+
+    # Create Mooncake config
+    config = create_mooncake_config(
+        host="127.0.0.1",
+        metadata_server="http://127.0.0.1:8080/metadata",
+        master="127.0.0.1:50051"
+    )
+
+    print(f"Created config with {len(config.connectors)} connectors")
+    for edge, connector_spec in config.connectors.items():
+        print(f"  {edge}: {connector_spec.name}")
+
+    return config
+
+def test_omni_llm_initialization():
+    """Test OmniLLM with connector config."""
+    print("\nTesting OmniLLM initialization with connectors...")
+
+    try:
+        from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+        # Create connector config
+        connector_config = {
+            "connectors": {
+                "0_to_1": {
+                    "name": "InMemoryConnector",  # Use InMemory for testing
+                    "extra": {}
+                }
+            }
+        }
+
+        # Test connector initialization logic
+        omni = OmniLLM.__new__(OmniLLM)  # Create without calling __init__
+        omni.omni_transfer_config = connector_config
+        omni.connectors = {}
+        omni._initialize_connectors()
+
+        # Check if connector was created
+        if ("0", "1") in omni.connectors:
+            print(f"✅ Connector created: {type(omni.connectors[('0', '1')]).__name__}")
+            return True
+        else:
+            print("❌ Connector not created")
+            return False
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+if __name__ == "__main__":
+    print("OmniConnector Integration Test")
+    print("=" * 40)
+
+    # Test config creation
+    config = test_connector_config()
+
+    # Test OmniLLM integration
+    success = test_omni_llm_initialization()
+
+    if success:
+        print("\n✅ Integration test passed!")
+        print("Connector system is ready for use with OmniLLM")
+    else:
+        print("\n❌ Integration test failed!")

vllm_omni/distributed/connectors/connector_tests/test_epdg_pipeline.py

@@ -0,0 +1,332 @@
+#!/usr/bin/env python3
+"""
+Test EPD (Encode-Prefill-Decode) 3-stage pipeline with connectors.
+Tests data flow between stages using configured connectors.
+"""
+
+import os
+import sys
+from typing import Any, Dict, Optional
+from dataclasses import dataclass
+import uuid
+
+try:
+    import yaml
+except ImportError:
+    print("PyYAML not available, using basic YAML parsing")
+    yaml = None
+
+# Add the connectors package to path
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+sys.path.insert(0, parent_dir)
+
+# Import from parent package
+try:
+    from connectors.config import ConnectorSpec, OmniTransferConfig
+    from connectors.base import OmniConnectorBase, InMemoryOmniConnector
+    from connectors.factory import OmniConnectorFactory
+    from connectors.utils import load_omni_transfer_config
+except ImportError:
+    # Fallback for direct execution
+    sys.path.insert(0, current_dir)
+    from config import ConnectorSpec, OmniTransferConfig
+    from base import OmniConnectorBase, InMemoryOmniConnector
+    from factory import OmniConnectorFactory
+    from utils import load_omni_transfer_config
+
+
+@dataclass
+class StageResult:
+    """Result from a stage processing."""
+    stage_name: str
+    request_id: str
+    output_data: Any
+    metadata: Dict[str, Any]
+
+
+class EncodeStage:
+    """Simple Encode stage for testing."""
+
+    def __init__(self, stage_id: int, output_connector: Optional[OmniConnectorBase] = None):
+        self.stage_id = stage_id
+        self.stage_name = "encode"
+        self.output_connector = output_connector
+
+    def process(self, request_id: str, input_data: Any) -> StageResult:
+        """Process input data and optionally send to next stage."""
+        print(f"[{self.stage_name}] Processing request {request_id}")
+
+        # Simple encoding simulation
+        encoded_data = {
+            "original_text": input_data,
+            "encoded_tokens": [ord(c) for c in str(input_data)][:10],  # Simple tokenization
+            "embeddings": [0.1 * i for i in range(10)],  # Mock embeddings
+            "stage": self.stage_name
+        }
+
+        result = StageResult(
+            stage_name=self.stage_name,
+            request_id=request_id,
+            output_data=encoded_data,
+            metadata={"tokens_count": len(encoded_data["encoded_tokens"])}
+        )
+
+        # Send to next stage if connector is available
+        if self.output_connector:
+            success = self.output_connector.put(
+                str(self.stage_id), "1", request_id, encoded_data
+            )
+            print(f"[{self.stage_name}] Sent data to stage 1 via connector: {success}")
+
+        return result
+
+
+class PrefillStage:
+    """Simple Prefill stage for testing."""
+
+    def __init__(self, stage_id: int,
+                 input_connector: Optional[OmniConnectorBase] = None,
+                 output_connector: Optional[OmniConnectorBase] = None):
+        self.stage_id = stage_id
+        self.stage_name = "prefill"
+        self.input_connector = input_connector
+        self.output_connector = output_connector
+
+    def process(self, request_id: str) -> Optional[StageResult]:
+        """Process data from previous stage."""
+        print(f"[{self.stage_name}] Processing request {request_id}")
+
+        # Try to get data from input connector
+        input_data = None
+        if self.input_connector:
+            input_data = self.input_connector.get("0", str(self.stage_id), request_id)
+
+        if input_data is None:
+            print(f"[{self.stage_name}] No input data received for {request_id}")
+            return None
+
+        print(f"[{self.stage_name}] Received data from encode")
+
+        # Simple prefill simulation
+        prefills_data = {
+            **input_data,
+            "kv_cache": {"layer_0": [0.2 * i for i in range(20)]},  # Mock KV cache
+            "stage": self.stage_name
+        }
+
+        result = StageResult(
+            stage_name=self.stage_name,
+            request_id=request_id,
+            output_data=prefills_data,
+            metadata={"kv_size": len(prefills_data["kv_cache"])}
+        )
+
+        # Send to next stage if connector is available
+        if self.output_connector:
+            success = self.output_connector.put(
+                str(self.stage_id), "2", request_id, prefills_data
+            )
+            print(f"[{self.stage_name}] Sent data to stage 2 via connector: {success}")
+
+        return result
+
+
+class DecodeStage:
+    """Simple Decode stage for testing (final stage in 3-stage pipeline)."""
+
+    def __init__(self, stage_id: int,
+                 input_connector: Optional[OmniConnectorBase] = None):
+        self.stage_id = stage_id
+        self.stage_name = "decode"
+        self.input_connector = input_connector
+
+    def process(self, request_id: str) -> Optional[StageResult]:
+        """Process data from previous stage (final output)."""
+        print(f"[{self.stage_name}] Processing request {request_id}")
+
+        # Try to get data from input connector
+        input_data = None
+        if self.input_connector:
+            input_data = self.input_connector.get("1", str(self.stage_id), request_id)
+
+        if input_data is None:
+            print(f"[{self.stage_name}] No input data received for {request_id}")
+            return None
+
+        print(f"[{self.stage_name}] Received data from prefill")
+
+        # Simple decode simulation (final output)
+        decoded_data = {
+            **input_data,
+            "hidden_states": [0.3 * i for i in range(15)],  # Mock hidden states
+            "final_output": f"Decoded text from: {input_data.get('original_text', 'unknown')}",
+            "stage": self.stage_name
+        }
+
+        result = StageResult(
+            stage_name=self.stage_name,
+            request_id=request_id,
+            output_data=decoded_data,
+            metadata={"hidden_size": len(decoded_data["hidden_states"]), "final": True}
+        )
+
+        return result
+
+
+class GeneratorStage:
+    """Simple Generator stage for testing (removed for 3-stage pipeline)."""
+    pass  # Not used in current 3-stage setup
+
+
+def load_config_from_yaml(yaml_path: str) -> OmniTransferConfig:
+    """Load connector configuration from YAML file."""
+    if yaml is None:
+        # Fallback: create a simple test config
+        print("Using fallback InMemoryConnector config for testing")
+        return create_test_config()
+
+    with open(yaml_path, 'r', encoding='utf-8') as f:
+        config_dict = yaml.safe_load(f)
+
+    try:
+        config = load_omni_transfer_config(config_dict=config_dict)
+        print(f"Successfully loaded config from YAML with {len(config.connectors)} connectors")
+        return config
+    except Exception as e:
+        print(f"Failed to load YAML config: {e}, using fallback config")
+        return create_test_config()
+
+
+def create_test_config() -> OmniTransferConfig:
+    """Create a test configuration with InMemoryConnectors."""
+    from .config import ConnectorSpec
+
+    connectors = {}
+    # Create connectors for EPD (3-stage) pipeline: 0->1->2
+    pipeline_stages = [("0", "1"), ("1", "2")]
+
+    for from_stage, to_stage in pipeline_stages:
+        connectors[(from_stage, to_stage)] = ConnectorSpec(
+            name="InMemoryConnector",
+            extra={}
+        )
+
+    return OmniTransferConfig(connectors=connectors)
+
+
+def create_stages_with_connectors(config: OmniTransferConfig) -> Dict[str, Any]:
+    """Create EPD (3-stage) pipeline with appropriate connectors."""
+
+    # Create connector instances based on config
+    connectors = {}
+    for edge_key, connector_spec in config.connectors.items():
+        from_stage, to_stage = edge_key
+        connectors[edge_key] = OmniConnectorFactory.create_connector(connector_spec)
+
+    # Create stages with connectors (3-stage pipeline: 0->1->2)
+    stages = {}
+
+    # Encode stage (stage 0)
+    encode_output_connector = connectors.get(("0", "1"))  # to prefill
+    stages["encode"] = EncodeStage(0, encode_output_connector)
+
+    # Prefill stage (stage 1)
+    prefill_input_connector = connectors.get(("0", "1"))   # from encode
+    prefill_output_connector = connectors.get(("1", "2"))  # to decode
+    stages["prefill"] = PrefillStage(1, prefill_input_connector, prefill_output_connector)
+
+    # Decode stage (stage 2) - final stage
+    decode_input_connector = connectors.get(("1", "2"))    # from prefill
+    stages["decode"] = DecodeStage(2, decode_input_connector)
+
+    return stages
+
+
+def test_pipeline_with_config(yaml_path: str):
+    """Test the EPDG pipeline with connector configuration."""
+
+    print("🔧 Loading configuration from YAML...")
+    try:
+        config = load_config_from_yaml(yaml_path)
+        print(f"✓ Loaded config with {len(config.connectors)} connector configurations")
+    except Exception as e:
+        print(f"✗ Failed to load config: {e}")
+        return False
+
+    print("\n🏗️  Creating EPDG stages with connectors...")
+    try:
+        stages = create_stages_with_connectors(config)
+        print("✓ Created all stages with connectors")
+    except Exception as e:
+        print(f"✗ Failed to create stages: {e}")
+        return False
+
+    print("\n🚀 Testing data flow through pipeline...")
+
+    # Test request
+    request_id = f"test_req_{uuid.uuid4()}"
+    input_text = "Hello, EPDG pipeline!"
+
+    try:
+        # Stage 1: Encode
+        print(f"\n--- Stage 1: Encode ---")
+        encode_result = stages["encode"].process(request_id, input_text)
+        print(f"✓ Encode completed: {encode_result.metadata}")
+
+        # Stage 2: Prefill
+        print(f"\n--- Stage 2: Prefill ---")
+        prefill_result = stages["prefill"].process(request_id)
+        if prefill_result:
+            print(f"✓ Prefill completed: {prefill_result.metadata}")
+        else:
+            print("✗ Prefill failed - no input data")
+            return False
+
+        # Stage 3: Decode
+        print(f"\n--- Stage 3: Decode ---")
+        decode_result = stages["decode"].process(request_id)
+        if decode_result:
+            print(f"✓ Decode completed: {decode_result.metadata}")
+        else:
+            print("✗ Decode failed - no input data")
+            return False
+
+        # Decode stage produces final output in 3-stage pipeline
+
+    except Exception as e:
+        print(f"✗ Pipeline test failed: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+    print("\n🎉 Pipeline test completed successfully!")
+    return True
+
+
+def main():
+    """Main test function."""
+    # Path to the YAML config file
+    yaml_path = "/workspace/omni/vllm-omni/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml"
+
+    if not os.path.exists(yaml_path):
+        print(f"✗ YAML config file not found: {yaml_path}")
+        return False
+
+    print("Testing EPDG Pipeline with Connectors")
+    print("=" * 50)
+    print(f"Config file: {yaml_path}")
+
+    success = test_pipeline_with_config(yaml_path)
+
+    if success:
+        print("\n✅ All tests passed!")
+    else:
+        print("\n❌ Tests failed!")
+
+    return success
+
+
+if __name__ == "__main__":
+    success = main()
+    sys.exit(0 if success else 1)

vllm_omni/distributed/connectors/factory.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+try:
+    from .base import OmniConnectorBase
+    from .config import ConnectorSpec
+except ImportError:
+    # Fallback for direct execution
+    import sys
+    import os
+    sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+    from connectors.base import OmniConnectorBase
+    from connectors.config import ConnectorSpec
+
+logger = logging.getLogger(__name__)
+
+
+class OmniConnectorFactory:
+    """Factory for creating OmniConnectors."""
+
+    _registry: dict[str, Callable[[dict[str, Any]], OmniConnectorBase]] = {}
+
+    @classmethod
+    def register_connector(
+        cls,
+        name: str,
+        constructor: Callable[[dict[str, Any]], OmniConnectorBase]
+    ) -> None:
+        """Register a connector constructor."""
+        if name in cls._registry:
+            raise ValueError(f"Connector '{name}' is already registered.")
+        cls._registry[name] = constructor
+        logger.debug(f"Registered connector: {name}")
+
+    @classmethod
+    def create_connector(cls, spec: ConnectorSpec) -> OmniConnectorBase:
+        """Create a connector from specification."""
+        if spec.name not in cls._registry:
+            raise ValueError(f"Unknown connector: {spec.name}. Available: {list(cls._registry.keys())}")
+
+        constructor = cls._registry[spec.name]
+        try:
+            connector = constructor(spec.extra)
+            logger.info(f"Created connector: {spec.name}")
+            return connector
+        except Exception as e:
+            logger.error(f"Failed to create connector {spec.name}: {e}")
+            raise
+
+    @classmethod
+    def list_registered_connectors(cls) -> list[str]:
+        """List all registered connector names."""
+        return list(cls._registry.keys())
+
+
+# Register built-in connectors with lazy imports
+def _create_mooncake_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .mooncake_connector import MooncakeConnector
+    except ImportError:
+        # Fallback import
+        import sys
+        import os
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from connectors.mooncake_connector import MooncakeConnector
+    return MooncakeConnector(config)
+
+
+def _create_in_memory_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .base import InMemoryOmniConnector
+    except ImportError:
+        # Fallback import
+        import sys
+        import os
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from connectors.base import InMemoryOmniConnector
+    return InMemoryOmniConnector()
+
+
+# Register connectors
+OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_connector)
+OmniConnectorFactory.register_connector("InMemoryConnector", _create_in_memory_connector)

vllm_omni/distributed/connectors/mooncake_connector.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import io
+import logging
+import time
+from typing import Any, Optional
+
+import torch
+
+from .base import OmniConnectorBase
+
+logger = logging.getLogger(__name__)
+
+try:
+    from mooncake.store import MooncakeDistributedStore, ReplicateConfig
+except ImportError:
+    try:
+        from mooncake import MooncakeDistributedStore, ReplicateConfig
+    except ImportError:
+        logger.warning("Mooncake not available, MooncakeOmniConnector will not work")
+        MooncakeDistributedStore = None
+        ReplicateConfig = None
+
+
+class MooncakeConnector(OmniConnectorBase):
+    """Mooncake-based distributed connector for OmniConnector."""
+
+    def __init__(self, config: dict[str, Any]):
+        if MooncakeDistributedStore is None:
+            raise ImportError("Mooncake not available")
+
+        self.config = config
+        self.host = config.get("host", "127.0.0.1")
+        self.metadata = config.get("metadata_server", "http://127.0.0.1:8080/metadata")
+        self.master = config.get("master", "127.0.0.1:50051")
+        self.segment = config.get("segment", 512 * 1024 * 1024)  # 512MB
+        self.localbuf = config.get("localbuf", 64 * 1024 * 1024)  # 64MB
+        self.proto = config.get("proto", "tcp")
+        self.rdma = config.get("rdma", "")
+
+        self.store: Optional[MooncakeDistributedStore] = None
+        self.pin: Optional[ReplicateConfig] = None
+
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "bytes_transferred": 0,
+            "errors": 0,
+            "timeouts": 0,
+        }
+
+        self._init_store()
+
+    def _make_key(self, rid: str, from_stage: str, to_stage: str) -> str:
+        """Generate store key for request between stages."""
+        return f"{rid}/{from_stage}_to_{to_stage}"
+
+    def _init_store(self):
+        """Initialize Mooncake store."""
+        try:
+            self.store = MooncakeDistributedStore()
+            rc = self.store.setup(
+                self.host, self.metadata, self.segment,
+                self.localbuf, self.proto, self.rdma, self.master
+            )
+            if rc != 0:
+                raise RuntimeError(f"Mooncake setup failed: {rc}")
+
+            self.pin = ReplicateConfig()
+            self.pin.with_soft_pin = True
+            logger.info("MooncakeConnector initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to initialize Mooncake store: {e}")
+            raise
+
+    # Use base class serialization methods for consistency
+
+    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
+        if not self.store:
+            logger.error("Store not initialized")
+            return False, 0
+
+        try:
+            serialized_data = self.serialize_obj(data)
+            key = self._make_key(request_id, from_stage, to_stage)
+            self.store.put(key, serialized_data, self.pin)
+
+            self._metrics["puts"] += 1
+            self._metrics["bytes_transferred"] += len(serialized_data)
+
+            logger.debug(f"MooncakeConnector: stored {key}, {len(serialized_data)} bytes")
+            print(f"========= mooncake put {from_stage}-{to_stage}-{request_id}, with size {len(serialized_data)}")
+            return True, len(serialized_data)
+
+        except Exception as e:
+            self._metrics["errors"] += 1
+            logger.error(f"MooncakeConnector put failed: {e}")
+            return False, 0
+
+    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
+        if not self.store:
+            logger.error("Store not initialized")
+            return None
+
+        retries = 20
+        sleep_s = 0.05
+        key = self._make_key(request_id, from_stage, to_stage)
+
+        for attempt in range(retries):
+            try:
+                raw_data = self.store.get(key)
+
+                if raw_data:
+                    data = self.deserialize_obj(raw_data)
+                    self._metrics["gets"] += 1
+                    print(f"========= mooncake get {from_stage}-{to_stage}-{request_id}, with size {len(raw_data)}")
+                    return data
+
+            except Exception as e:
+                logger.debug(f"MooncakeConnector get attempt {attempt} failed: {e}")
+
+            if attempt < retries - 1:
+                time.sleep(sleep_s)
+
+        self._metrics["timeouts"] += 1
+        logger.warning(f"MooncakeConnector: timeout waiting for {key}")
+        return None
+
+    def cleanup(self, request_id: str) -> None:
+        if not self.store:
+            return
+
+        # Note: Mooncake doesn't have explicit delete, data will be garbage collected
+        # We could implement a cleanup mechanism by storing deletion markers
+        logger.debug(f"MooncakeConnector: cleanup requested for {request_id} (no-op)")
+
+    def health(self) -> dict[str, Any]:
+        if not self.store:
+            return {"status": "unhealthy", "error": "Store not initialized"}
+
+        return {
+            "status": "healthy",
+            "host": self.host,
+            "metadata_server": self.metadata,
+            "master": self.master,
+            **self._metrics,
+        }
+
+    def close(self):
+        """Clean shutdown."""
+        if self.store:
+            try:
+                self.store.close()
+                self.store = None
+                logger.info("MooncakeConnector closed")
+            except Exception as e:
+                logger.error(f"Error closing Mooncake store: {e}")

vllm_omni/distributed/connectors/utils.py

@@ -0,0 +1,314 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Utilities for OmniConnector configuration and validation."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple, Union
+
+from vllm_omni.distributed.connectors.config import OmniTransferConfig, ConnectorSpec
+from vllm_omni.distributed.connectors.factory import OmniConnectorFactory
+from vllm_omni.distributed.connectors.base import OmniConnectorBase
+
+logger = logging.getLogger(__name__)
+
+
+def initialize_connectors_from_config(
+    config_path: Optional[Union[str, Path]] = None
+) -> tuple[Optional[OmniTransferConfig], dict[tuple[str, str], OmniConnectorBase]]:
+    """
+    Initialize connectors from configuration file.
+
+    Returns:
+        tuple: (OmniTransferConfig, dict of {(from, to): connector_instance})
+    """
+    try:
+        transfer_config = load_omni_transfer_config(config_path)
+    except Exception as e:
+        logger.warning(f"Failed to load OmniTransferConfig from {config_path}: {e}")
+        return None, {}
+
+    if not transfer_config:
+        logger.info("No OmniTransferConfig provided, using default IPC")
+        return None, {}
+
+    # 使用统一的连接器创建逻辑
+    connectors = create_connectors_from_config(transfer_config.connectors)
+    return transfer_config, connectors
+
+
+def create_connectors_from_config(
+    connectors_config: dict[tuple[str, str], ConnectorSpec]
+) -> dict[tuple[str, str], OmniConnectorBase]:
+    """通用连接器创建逻辑，从配置字典创建连接器实例。"""
+    connectors = {}
+    try:
+        for edge_key, connector_spec in connectors_config.items():
+            connector = OmniConnectorFactory.create_connector(connector_spec)
+            connectors[edge_key] = connector
+            logger.info(f"Created connector for {edge_key[0]} -> {edge_key[1]}: {type(connector).__name__}")
+    except Exception as e:
+        logger.error(f"Failed to initialize connectors: {e}")
+
+    return connectors
+
+
+def get_connectors_config_for_stage(
+    transfer_config: Optional[OmniTransferConfig],
+    stage_id: Union[str, int]
+) -> dict[str, Any]:
+    """
+    Extract connector configurations relevant for a specific stage worker.
+    
+    Returns a dict compatible with worker initialization:
+    {
+        "from_stage_X": {
+            "spec": {
+                "name": "ConnectorName",
+                "extra": {...}
+            }
+        },
+        ...
+    }
+    """
+    if not transfer_config:
+        return {}
+        
+    stage_connectors_config = {}
+    target_stage = str(stage_id)
+    
+    # Iterate through all configured edges
+    for (from_stage, to_stage), spec in transfer_config.connectors.items():
+        # We only care about incoming edges for the worker process
+        # (Worker needs to create connectors to receive data)
+        if to_stage == target_stage:
+            stage_connectors_config[f"from_stage_{from_stage}"] = {
+                "spec": {
+                    "name": spec.name,
+                    "extra": spec.extra
+                }
+            }
+            
+    return stage_connectors_config
+
+
+def load_omni_transfer_config(
+    config_path: Optional[Union[str, Path]] = None,
+    config_dict: Optional[dict[str, Any]] = None
+) -> Optional[OmniTransferConfig]:
+    """Load OmniTransferConfig from file or dict."""
+    if config_path is None and config_dict is None:
+        return None
+
+    if config_path is not None:
+        config_path = Path(config_path)
+        if not config_path.exists():
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+
+        with open(config_path, 'r', encoding='utf-8') as f:
+            if config_path.suffix.lower() == '.json':
+                config_dict = json.load(f)
+            elif config_path.suffix.lower() in ['.yaml', '.yml']:
+                try:
+                    import yaml
+                    config_dict = yaml.safe_load(f)
+                except ImportError:
+                    raise ImportError("PyYAML required for YAML config files")
+            else:
+                raise ValueError(f"Unsupported config file format: {config_path.suffix}")
+
+    if config_dict is None:
+        return None
+
+    # Parse connectors
+    connectors = {}
+    runtime_config = config_dict.get('runtime', {})
+
+    # Parse global connectors (from runtime.connectors)
+    global_connectors = runtime_config.get('connectors', {})
+    for conn_name, conn_config in global_connectors.items():
+        connector = ConnectorSpec(
+            name=conn_config['name'],
+            extra=conn_config.get('extra', {})
+        )
+        # Store globally defined connectors for reference
+        # These will be resolved to specific edges later
+
+    # Parse stage-level connectors
+    for stage_config in config_dict.get('stage_args', []):
+        stage_id = str(stage_config['stage_id'])
+
+        # Input connectors
+        for input_key, conn_ref in stage_config.get('input_connectors', {}).items():
+            if isinstance(conn_ref, str):
+                # Reference to global connector
+                if conn_ref in global_connectors:
+                    conn_config = global_connectors[conn_ref]
+                    connector = ConnectorSpec(
+                        name=conn_config['name'],
+                        extra=conn_config.get('extra', {})
+                    )
+                else:
+                    raise ValueError(f"Undefined connector reference: {conn_ref}")
+            else:
+                # Inline connector definition
+                connector = ConnectorSpec(
+                    name=conn_ref['name'],
+                    extra=conn_ref.get('extra', {})
+                )
+
+            # Parse from_stage from key (e.g., "from_stage_0" -> "0")
+            from_stage = input_key.replace("from_stage_", "")
+            edge_key = (from_stage, stage_id)
+            connectors[edge_key] = connector
+
+        # Output connectors
+        for output_key, conn_ref in stage_config.get('output_connectors', {}).items():
+            if isinstance(conn_ref, str):
+                # Reference to global connector
+                if conn_ref in global_connectors:
+                    conn_config = global_connectors[conn_ref]
+                    connector = ConnectorSpec(
+                        name=conn_config['name'],
+                        extra=conn_config.get('extra', {})
+                    )
+                else:
+                    raise ValueError(f"Undefined connector reference: {conn_ref}")
+            else:
+                # Inline connector definition
+                connector = ConnectorSpec(
+                    name=conn_ref['name'],
+                    extra=conn_ref.get('extra', {})
+                )
+
+            # Parse to_stage from key (e.g., "to_stage_1" -> "1")
+            to_stage = output_key.replace("to_stage_", "")
+            edge_key = (stage_id, to_stage)
+            connectors[edge_key] = connector
+
+    config = OmniTransferConfig(connectors=connectors)
+
+    logger.info(f"Loaded OmniTransferConfig with {len(connectors)} connector configurations")
+    return config
+
+
+def create_simple_config(
+    connector_name: str = "InMemoryConnector",
+    connector_config: Optional[dict[str, Any]] = None
+) -> OmniTransferConfig:
+    """Create a simple OmniTransferConfig for testing."""
+    if connector_config is None:
+        connector_config = {}
+
+    # Create connectors for typical EPDG pipeline
+    connectors = {}
+    pipeline_edges = [
+        ("0", "1"),  # thinker -> talker
+        ("1", "2"),  # talker -> code2wav
+    ]
+
+    for from_stage, to_stage in pipeline_edges:
+        connectors[(from_stage, to_stage)] = ConnectorSpec(
+            name=connector_name,
+            extra=connector_config
+        )
+
+    return OmniTransferConfig(connectors=connectors)
+
+
+def create_mooncake_config(
+    host: str = "127.0.0.1",
+    metadata_server: str = "http://127.0.0.1:8080/metadata",
+    master: str = "127.0.0.1:50051",
+    segment: int = 512 * 1024 * 1024,
+    localbuf: int = 64 * 1024 * 1024,
+) -> OmniTransferConfig:
+    """Create OmniTransferConfig for Mooncake-based setup."""
+    connector_config = {
+        "host": host,
+        "metadata_server": metadata_server,
+        "master": master,
+        "segment": segment,
+        "localbuf": localbuf,
+    }
+
+    return create_simple_config("MooncakeConnector", connector_config)
+
+
+# High-level management functions (moved from manager.py)
+
+def initialize_orchestrator_connectors(
+    config_path: Optional[str],
+) -> Tuple[Optional[OmniTransferConfig], Dict[tuple[str, str], OmniConnectorBase]]:
+    """Initialize connectors shared at orchestrator level."""
+    try:
+        transfer_config, connectors = initialize_connectors_from_config(config_path)
+        return transfer_config, connectors
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.error("Unexpected error initializing connectors: %s", exc)
+        return None, {}
+
+
+def get_stage_connector_config(
+    transfer_config: Optional[OmniTransferConfig],
+    stage_id: int,
+) -> dict[str, Any]:
+    """Return the serialized connector config payload for a specific stage."""
+    if transfer_config is None:
+        return {}
+
+    try:
+        return get_connectors_config_for_stage(transfer_config, stage_id)
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.warning(
+            "Failed to build connector config for stage %s: %s. Using IPC fallback.",
+            stage_id,
+            exc,
+        )
+        return {}
+
+
+def build_stage_connectors(
+    stage_id: int,
+    connectors_config: dict[str, Any],
+) -> Optional[dict[tuple[str, str], Any]]:
+    """Instantiate OmniConnectors for a stage based on config."""
+    if not connectors_config:
+        return {}
+
+    logger.info(
+        "[Stage-%s] Initializing OmniConnectors with config keys: %s",
+        stage_id,
+        list(connectors_config.keys()),
+    )
+
+    from .factory import OmniConnectorFactory
+    from .config import ConnectorSpec
+
+    connectors: dict[tuple[str, str], Any] = {}
+    try:
+        # 将字典格式的配置转换为ConnectorSpec对象
+        stage_connector_specs = {}
+        for input_key, config in connectors_config.items():
+            if not input_key.startswith("from_stage_"):
+                continue
+
+            from_stage = input_key.replace("from_stage_", "")
+            spec_dict = config.get("spec", {})
+            if not spec_dict:
+                continue
+
+            connector_spec = ConnectorSpec(
+                name=spec_dict.get("name", "InMemoryConnector"),
+                extra=spec_dict.get("extra", {}),
+            )
+            stage_connector_specs[(str(from_stage), str(stage_id))] = connector_spec
+
+        # 使用统一的连接器创建逻辑
+        connectors = create_connectors_from_config(stage_connector_specs)
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc)
+        return None
+
+    return connectors

vllm_omni/entrypoints/async_omni.py

@@ -51,7 +51,7 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
-class AsyncOmniLLM(EngineClient):
+class AsyncOmni(EngineClient):
     """Async entry point for vLLM-Omni inference.
 
     This class provides an asynchronous interface for running multi-modal
@@ -75,7 +75,7 @@ class AsyncOmniLLM(EngineClient):
         **kwargs: Additional keyword arguments passed to stage engines
 
     Example:
-        >>> async_llm = AsyncOmniLLM(model="Qwen/Qwen2.5-Omni-7B")
+        >>> async_llm = AsyncOmni(model="Qwen/Qwen2.5-Omni-7B")
         >>> async for output in async_llm.generate(
         ...     prompt="Hello",
         ...     request_id="req-1",
@@ -94,7 +94,7 @@ class AsyncOmniLLM(EngineClient):
         shm_threshold_bytes: int = 65536,
         batch_timeout: int = 10,
         init_timeout: int = 60000,
-        **kwargs,
+        **kwargs: Any,
     ):
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
@@ -175,7 +175,7 @@ class AsyncOmniLLM(EngineClient):
 
         Sends shutdown signals to all stage input queues and stops
         all stage worker processes. This method should be called
-        when done using the AsyncOmniLLM instance.
+        when done using the AsyncOmni instance.
         """
         for q in self._stage_in_queues:
             try:
@@ -193,7 +193,7 @@ class AsyncOmniLLM(EngineClient):
                 logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
 
     def __del__(self) -> None:  # best-effort
-        print("[AsyncOmniLLM] __del__ close()", flush=True)
+        print("[AsyncOmni] __del__ close()", flush=True)
         raise Exception("test")
         try:
             self.close()
@@ -611,13 +611,13 @@ class AsyncOmniLLM(EngineClient):
         **kwargs,
     ):
         """Generate outputs for a request from a pooling model."""
-        raise NotImplementedError("encode() is not implemented for AsyncOmniLLM")
+        raise NotImplementedError("encode() is not implemented for AsyncOmni")
 
     async def start_profile(self) -> None:
-        raise NotImplementedError("start_profile() is not implemented for AsyncOmniLLM")
+        raise NotImplementedError("start_profile() is not implemented for AsyncOmni")
 
     async def stop_profile(self) -> None:
-        raise NotImplementedError("stop_profile() is not implemented for AsyncOmniLLM")
+        raise NotImplementedError("stop_profile() is not implemented for AsyncOmni")
 
 
 class AsyncOmniStageLLM(AsyncLLM):

vllm_omni/entrypoints/omni_llm.py

@@ -1,6 +1,8 @@
 import multiprocessing as mp
 import os
+import sys
 import time
+import uuid
 from collections.abc import Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Optional, Union
@@ -20,6 +22,11 @@ from vllm.utils import Counter
 from vllm.v1.engine.llm_engine import LLMEngine
 
 # Internal imports (our code)
+from vllm_omni.distributed.connectors.adapter import try_send_via_connector
+from vllm_omni.distributed.connectors import (
+    initialize_orchestrator_connectors,
+    get_stage_connector_config,
+)
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
@@ -31,9 +38,14 @@ from vllm_omni.entrypoints.log_utils import (
 )
 from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
+from vllm_omni.entrypoints.stage_utils import _to_dict
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
+from vllm_omni.entrypoints.utils import (
+    load_stage_configs_from_model,
+    load_stage_configs_from_yaml,
+    resolve_model_config_path,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -84,14 +96,23 @@ class OmniLLM:
         init_timeout: int = 300,
         **kwargs: Any,
     ):
+        self.worker_backend = kwargs.get("worker_backend", "process")
+        self.ray_address = kwargs.get("ray_address", None)
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
+        self._ray_pg = None
+
         # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
         if stage_configs_path is None:
-            self.stage_configs = load_stage_configs_from_model(model)
+            self.config_path = resolve_model_config_path(model)
+            self.stage_configs = load_stage_configs_from_yaml(self.config_path)
         else:
+            self.config_path = stage_configs_path
             self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
 
+        # Initialize connectors
+        self.omni_transfer_config,self.connectors = initialize_orchestrator_connectors(self.config_path)
+
         # Optional file handler for orchestrator
         self._log_file = log_file
 
@@ -124,36 +145,140 @@ class OmniLLM:
         self.stage_list = [st for _, st in results]
         logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
-        self._ctx = mp.get_context("spawn")
+        if self.worker_backend == "ray":
+            try:
+                import ray
+                from ray.util.queue import Queue as RayQueue
+            except ImportError:
+                raise ImportError("ray is required for worker_backend='ray'")
+            self._queue_cls = lambda: RayQueue(maxsize=0)
+        else:
+            self._ctx = mp.get_context("spawn")
+            self._queue_cls = lambda: self._ctx.Queue(maxsize=0)
+
         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)
+        self._start_stages(model)
         # Wait for all stages to report readiness before seeding
         self._stages_ready: set[int] = set()
         self._wait_for_stages_ready(timeout=init_timeout)
 
-    def _start_stage_processes(self, model: str) -> None:
+    def _start_stages(self, model: str) -> None:
+        if self.worker_backend == "ray":
+            self._start_stage_actors_ray(model)
+        else:
+            self._start_stage_processes_mp(model)
+
+    def _start_stage_processes_mp(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)
+            in_q: mp.Queue = self._queue_cls()
-            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            out_q: mp.Queue = self._queue_cls()
             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)
+
+            # Build connectors config for this stage
+            stage_connectors_config = get_stage_connector_config(
+                self.omni_transfer_config,
+                stage_id,
+            )
+
             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,
+                connectors_config=stage_connectors_config,
             )
             logger.debug("[Orchestrator] Stage-%s process started", stage_id)
             time.sleep(self._init_sleep_seconds)
 
+    def _start_stage_actors_ray(self, model: str) -> None:
+        import ray
+        import os
+        from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+        # Initialize Ray Cluster
+        if not ray.is_initialized():
+             # Create a dummy parallel config just to reuse vLLM's init logic if needed
+             # or just call ray.init()
+             # Pass current PYTHONPATH to workers to ensure they can find vllm_omni
+             runtime_env = {
+                 "env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}
+             }
+             ray.init(address=self.ray_address, ignore_reinit_error=True, runtime_env=runtime_env)
+
+        # Create Placement Group
+        # We allocate 1 bundle per stage. Assume 1 GPU per stage for now.
+        # Must include CPU=1 because default Actor request includes 1 CPU.
+        bundles = [{"GPU": 1.0, "CPU": 1.0} for _ in self.stage_list]
+        self._ray_pg = ray.util.placement_group(bundles, strategy="SPREAD")
+        ray.get(self._ray_pg.ready())
+        logger.info("[Orchestrator] Ray Placement Group created")
+
+        # Define Actor Wrapper
+        @ray.remote(num_gpus=1)
+        class OmniStageRayWorker:
+            def run(self, func, *args, **kwargs):
+                return func(*args, **kwargs)
+
+        from vllm_omni.entrypoints.omni_stage import _stage_worker
+
+        for stage_id, stage in enumerate(self.stage_list):
+            in_q = self._queue_cls()
+            out_q = self._queue_cls()
+            self._stage_in_queues.append(in_q)
+            self._stage_out_queues.append(out_q)
+            stage.attach_queues(in_q, out_q)
+
+            stage_connectors_config = get_stage_connector_config(
+                self.omni_transfer_config,
+                stage_id,
+            )
+
+            # Prepare payload
+            # Reuse logic from OmniStage.init_stage_worker but do not start process
+            engine_args = _to_dict(stage.engine_args)
+            runtime_cfg = _to_dict(getattr(stage.stage_config, "runtime", {}))
+            stage_payload = {
+                "stage_id": stage.stage_id,
+                "engine_args": engine_args,
+                "runtime": runtime_cfg,
+                "shm_threshold_bytes": sys.maxsize, # Disable SHM for Ray mode by setting huge threshold
+                "connectors_config": stage_connectors_config or {},
+            }
+
+            # Start Actor
+            # Inject PYTHONPATH again for the specific actor to be safe
+            worker_actor = OmniStageRayWorker.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=self._ray_pg,
+                    placement_group_bundle_index=stage_id
+                ),
+                runtime_env={"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}}
+            ).remote()
+
+            # Start worker loop asynchronously
+            worker_actor.run.remote(
+                _stage_worker,
+                model=model,
+                stage_payload=stage_payload,
+                in_q=in_q,
+                out_q=out_q,
+                log_file=self._log_file,
+                batch_timeout=self.batch_timeout,
+            )
+            
+            stage._ray_actor = worker_actor
+            logger.debug("[Orchestrator] Stage-%s Ray Actor started", stage_id)
+            time.sleep(self._init_sleep_seconds)
+
     def close(self) -> None:
         """Close all stage processes and clean up resources.
 
@@ -171,10 +296,21 @@ class OmniLLM:
                 )
         for stage in self.stage_list:
             try:
-                stage.stop_stage_worker()
+                if self.worker_backend == "ray" and hasattr(stage, "_ray_actor"):
+                    import ray
+                    # Ray kill
+                    ray.kill(stage._ray_actor)
+                else:
+                    stage.stop_stage_worker()
             except Exception as e:
                 logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
 
+        if self.worker_backend == "ray" and self._ray_pg:
+            try:
+                ray.util.remove_placement_group(self._ray_pg)
+            except Exception as e:
+                logger.warning("[Orchestrator] Failed to remove placement group: %s", e)
+
     def __del__(self) -> None:  # best-effort
         try:
             self.close()
@@ -237,11 +373,14 @@ class OmniLLM:
         # Orchestrator keeps stage objects for input derivation
         num_stages = len(self.stage_list)
 
-        # Map from request_id to original prompt
+        # Generate globally unique request IDs and map them to original prompts
-        request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
+        request_ids: list[str] = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))]
+        request_id_to_prompt: dict[str, PromptType] = {
+            rid: p for rid, p in zip(request_ids, request_prompts)
+        }
 
         # Track per-request start time for end-to-end timing
-        _req_start_ts: dict[int, float] = {}
+        _req_start_ts: dict[str, float] = {}
         _wall_start_ts: float = time.time()
 
         # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
@@ -387,51 +526,90 @@ class OmniLLM:
                         )
                         continue
                     sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
-                    try:
+
-                        # Measure transfer size and time (encode + enqueue)
+                    # Check if we have a connector for this edge
-                        size_bytes = 0
+                    connector_key = (str(stage_id), str(next_stage_id))
+                    connector = self.connectors.get(connector_key)
+                    print(f"connectors: {self.connectors}")
+
+                    sent_via_connector = False
+                    if connector:
+                        sent_via_connector = try_send_via_connector(
+                            connector=connector,
+                            stage_id=stage_id,
+                            next_stage_id=next_stage_id,
+                            req_id=req_id,
+                            next_inputs=next_inputs,
+                            sampling_params=sp_next,
+                            original_prompt=request_id_to_prompt[req_id],
+                            next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
+                            metrics=metrics
+                        )
+
+                    if not sent_via_connector:
+                        # Use original queue mechanism
                         try:
-                            size_bytes = len(_set(next_inputs))
+                            if self.worker_backend == "ray":
-                        except Exception:
+                                # Ray mode: Skip SHM fallback, just send object via Queue (Ray handles serialization)
-                            size_bytes = 0
+                                # This prevents SHM access errors across nodes
-                        t0 = time.time()
+                                ipc_payload = {
-                        ipc_payload = _encode(
+                                    "request_id": req_id,
-                            next_inputs,
+                                    "engine_inputs": next_inputs,
-                            getattr(self, "_shm_threshold_bytes", 65536),
+                                    "sampling_params": sp_next,
-                            obj_key="engine_inputs",
+                                    "sent_ts": time.time(),
-                            shm_key="engine_inputs_shm",
+                                }
-                        )
+                                size_bytes = 0 # Approximation or calculate size if needed
-                        ipc_payload.update(
+                                try:
-                            {
+                                    # Approximate size for stats
-                                "request_id": req_id,
+                                    size_bytes = len(_set(next_inputs))
-                                "sampling_params": sp_next,
+                                except:
-                                "sent_ts": time.time(),
+                                    pass
-                            }
+                                t0 = time.time()
-                        )
+                                # Direct submit to Ray Queue
-                        self.stage_list[next_stage_id].submit(ipc_payload)
+                            else:
-                        t1 = time.time()
+                                size_bytes = 0
-                        tx_ms = (t1 - t0) * 1000.0
+                                try:
-                        metrics.on_forward(
+                                    size_bytes = len(_set(next_inputs))
-                            stage_id,
+                                except Exception:
-                            next_stage_id,
+                                    size_bytes = 0
-                            req_id,
+                                t0 = time.time()
-                            int(size_bytes),
+                                ipc_payload = _encode(
-                            float(tx_ms),
+                                    next_inputs,
-                            bool("engine_inputs_shm" in ipc_payload),
+                                    getattr(self, "_shm_threshold_bytes", 65536),
-                        )
+                                    obj_key="engine_inputs",
-                    except Exception as e:
+                                    shm_key="engine_inputs_shm",
-                        logger.warning(
+                                )
-                            "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
+                                ipc_payload.update(
-                            req_id,
+                                    {
-                            e,
+                                        "request_id": req_id,
-                        )
+                                        "sampling_params": sp_next,
-                        self.stage_list[next_stage_id].submit(
+                                        "sent_ts": time.time(),
-                            {
+                                    }
-                                "request_id": req_id,
+                                )
-                                "engine_inputs": next_inputs,
+                            self.stage_list[next_stage_id].submit(ipc_payload)
-                                "sampling_params": sp_next,
+                            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,
@@ -534,8 +712,9 @@ class OmniStageLLM(LLM):
         self,
         model: str,
         compilation_config: Optional[Union[int, dict[str, Any], CompilationConfig]] = None,
+        hf_overrides: Optional[dict[str, Any]] = None,
         structured_outputs_config: Optional[Union[dict[str, Any], StructuredOutputsConfig]] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         """LLM constructor."""
         if "disable_log_stats" not in kwargs:

vllm_omni/entrypoints/omni_stage.py

@@ -32,7 +32,7 @@ from vllm_omni.entrypoints.stage_utils import (
     _to_dict,
     maybe_dump_to_shm,
     maybe_load_from_ipc_with_metrics,
-    set_stage_gpu_devices,
+    set_stage_devices,
 )
 from vllm_omni.inputs.data import OmniTokensPrompt
 
@@ -163,6 +163,7 @@ class OmniStage:
         shm_threshold_bytes: int = 65536,
         ctx: Optional[mp.context.BaseContext] = None,
         batch_timeout: int = 10,
+        connectors_config: Optional[dict] = None,  # New parameter
     ) -> None:
         """Initialize and start the stage worker process.
 
@@ -192,6 +193,7 @@ class OmniStage:
             "engine_args": engine_args,
             "runtime": runtime_cfg,
             "shm_threshold_bytes": self._shm_threshold_bytes,
+            "connectors_config": connectors_config or {},
         }
         if is_async:
             self._proc = ctx.Process(
@@ -295,6 +297,8 @@ class OmniStage:
                 raise ValueError("engine_input_source is empty")
             source_stage_id = self.engine_input_source[0]
             source_outputs = stage_list[source_stage_id].engine_outputs
+            if not isinstance(prompt, list):
+                prompt = [prompt]
             multi_modal_data = {
                 source_output.request_id: p.get("multi_modal_data", None)
                 for source_output, p in zip(source_outputs, prompt)
@@ -338,6 +342,8 @@ def _stage_worker(
         log_stage_batch_stats,
         log_stage_running_avg,
     )
+    from vllm_omni.distributed.connectors.adapter import try_recv_via_connector
+    from vllm_omni.distributed.connectors.utils import build_stage_connectors
     from vllm_omni.entrypoints.omni_llm import OmniStageLLM
     # no inline JSONL/serialization imports; logging handled by utilities
 
@@ -345,6 +351,7 @@ def _stage_worker(
     engine_args = stage_payload.get("engine_args", {})
     runtime_cfg = stage_payload.get("runtime", {})
     shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
+    connectors_config = stage_payload.get("connectors_config", {})
 
     # Per-stage logger: clear inherited handlers to avoid broken parent streams
     try:
@@ -425,7 +432,10 @@ def _stage_worker(
 
     # Device mapping
     try:
-        set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
+        from vllm_omni.utils import detect_device_type
+
+        device_type = detect_device_type()
+        set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
     except Exception as e:
         _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
 
@@ -435,6 +445,18 @@ def _stage_worker(
     )
     stage_engine = OmniStageLLM(model=model, **engine_args)
     _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
+
+    # Initialize OmniConnectors if configured
+    connectors = {}
+    if connectors_config:
+        built_connectors = build_stage_connectors(
+            stage_id=stage_id,
+            connectors_config=connectors_config,
+        )
+        if built_connectors is None:
+            return
+        connectors = built_connectors
+
     # Signal readiness to orchestrator
     try:
         out_q.put({"type": "stage_ready", "stage_id": stage_id})
@@ -452,9 +474,9 @@ def _stage_worker(
         max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
         print(f"[Stage-{stage_id}] Max batch size: {max_batch_size}")
         batch_tasks: list[dict[str, Any]] = [task]
+        start_time = _time.time()
         if max_batch_size > 1:
             while len(batch_tasks) < max_batch_size:
-                start_time = _time.time()
                 if not in_q.empty():
                     extra = in_q.get_nowait()
                     if extra is None:
@@ -491,9 +513,21 @@ def _stage_worker(
                     _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))
+            # Resolve input data (via Connector or IPC)
-            _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
+            ein, _rx_metrics = try_recv_via_connector(
+                task=t,
+                connectors=connectors,
+                stage_id=stage_id
+            )
+
+            if ein is None or _rx_metrics is None:
+                ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+
+            if _rx_metrics:
+                _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)
@@ -666,7 +700,7 @@ async def _stage_worker_async(
     import logging as _logging
     import time as _time
 
-    from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM
+    from vllm_omni.entrypoints.async_omni import AsyncOmniStageLLM
     from vllm_omni.entrypoints.log_utils import (
         compute_and_log_stage_request_stats,
         count_tokens_from_outputs,
@@ -717,7 +751,10 @@ async def _stage_worker_async(
 
     # Device mapping
     try:
-        set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
+        from vllm_omni.utils import detect_device_type
+
+        device_type = detect_device_type()
+        set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
     except Exception as e:
         _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
 

vllm_omni/entrypoints/openai/api_server.py

@@ -38,7 +38,7 @@ from vllm.logger import init_logger
 from vllm.transformers_utils.tokenizer import MistralTokenizer
 from vllm.utils import decorate_logs
 
-from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
 
 logger = init_logger(__name__)
@@ -65,7 +65,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
     if log_config is not None:
         uvicorn_kwargs["log_config"] = log_config
 
-    async with build_async_omni_llm(
+    async with build_async_omni(
         args,
         client_config=client_config,
     ) as engine_client:
@@ -108,15 +108,15 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
 
 @asynccontextmanager
-async def build_async_omni_llm(
+async def build_async_omni(
     args: Namespace,
     *,
     disable_frontend_multiprocessing: Optional[bool] = None,
     client_config: Optional[dict[str, Any]] = None,
 ) -> AsyncIterator[EngineClient]:
-    """Build an AsyncOmniLLM instance from command-line arguments.
+    """Build an AsyncOmni instance from command-line arguments.
 
-    Creates an async context manager that yields an AsyncOmniLLM instance
+    Creates an async context manager that yields an AsyncOmni instance
     configured from the provided arguments. Handles forkserver setup if
     needed and ensures proper cleanup on exit.
 
@@ -127,7 +127,7 @@ async def build_async_omni_llm(
         client_config: Optional client configuration dictionary
 
     Yields:
-        EngineClient instance (AsyncOmniLLM) ready for use
+        EngineClient instance (AsyncOmni) ready for use
     """
     if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver":
         # The executor is expected to be mp.
@@ -138,25 +138,25 @@ async def build_async_omni_llm(
         forkserver.ensure_running()
         logger.debug("Forkserver setup complete!")
 
-    # Context manager to handle async_omni_llm lifecycle
+    # Context manager to handle async_omni lifecycle
     # Ensures everything is shutdown and cleaned up on error/exit
-    async with build_async_omni_llm_from_stage_config(
+    async with build_async_omni_from_stage_config(
         args,
         disable_frontend_multiprocessing=disable_frontend_multiprocessing,
-    ) as async_omni_llm:
+    ) as async_omni:
-        yield async_omni_llm
+        yield async_omni
 
 
 @asynccontextmanager
-async def build_async_omni_llm_from_stage_config(
+async def build_async_omni_from_stage_config(
     args: Namespace,
     *,
     disable_frontend_multiprocessing: bool = False,
     client_config: Optional[dict[str, Any]] = None,
 ) -> AsyncIterator[EngineClient]:
-    """Create AsyncOmniLLM from stage configuration.
+    """Create AsyncOmni from stage configuration.
 
-    Creates an AsyncOmniLLM instance either in-process or using multiprocess
+    Creates an AsyncOmni instance either in-process or using multiprocess
     RPC. Loads stage configurations from the model or from a specified path.
 
     Args:
@@ -166,7 +166,7 @@ async def build_async_omni_llm_from_stage_config(
         client_config: Optional client configuration dictionary
 
     Yields:
-        EngineClient instance (AsyncOmniLLM) ready for use
+        EngineClient instance (AsyncOmni) ready for use
 
     Note:
         Stage configurations are loaded from args.stage_configs_path if provided,
@@ -182,21 +182,21 @@ async def build_async_omni_llm_from_stage_config(
             "To disable frontend multiprocessing, set VLLM_USE_V1=0."
         )
 
-    async_omni_llm: Optional[EngineClient] = None
+    async_omni: Optional[EngineClient] = None
 
     try:
         if getattr(args, "stage_configs_path", None):
-            async_omni_llm = AsyncOmniLLM(model=args.model, stage_configs_path=args.stage_configs_path)
+            async_omni = AsyncOmni(model=args.model, stage_configs_path=args.stage_configs_path)
         else:
-            async_omni_llm = AsyncOmniLLM(model=args.model)
+            async_omni = AsyncOmni(model=args.model)
 
         # # Don't keep the dummy data in memory
         # await async_llm.reset_mm_cache()
 
-        yield async_omni_llm
+        yield async_omni
     finally:
-        if async_omni_llm:
+        if async_omni:
-            async_omni_llm.shutdown()
+            async_omni.shutdown()
 
 
 async def omni_init_app_state(
@@ -211,7 +211,7 @@ async def omni_init_app_state(
     and other server configuration needed for handling API requests.
 
     Args:
-        engine_client: Engine client instance (AsyncOmniLLM)
+        engine_client: Engine client instance (AsyncOmni)
         vllm_config: vLLM configuration object
         state: FastAPI application state object to initialize
         args: Parsed command-line arguments

vllm_omni/entrypoints/stage_utils.py

@@ -13,100 +13,143 @@ from omegaconf import OmegaConf
 logger = logging.getLogger(__name__)
 
 
-def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
+def set_stage_devices(
-    """Configure per-stage CUDA visibility and current device.
+    stage_id: int,
+    devices: str | int | None,
+    device_type: str | None = None,
+) -> None:
+    """Configure per-stage device visibility and current device (CUDA or NPU).
 
-    Behavior
+    This function sets environment variables that control which devices are visible
-    - Comma-separated string (e.g. "2,5,7"): set CUDA_VISIBLE_DEVICES exactly
+    to the process, and sets the current device. It must be called BEFORE worker
-      to this list; logical index 0 is used as current device.
+    initialization so that workers see the correct devices.
-    - 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
+    Args:
-      CUDA_VISIBLE_DEVICES to this single device.
+        stage_id: Stage identifier for logging
-    - None/"cpu": keep default visibility.
+        devices: Device specification:
-    - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string.
+            - Comma-separated string (e.g. "2,5,7"): set device visibility env var
+              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 device visibility mapping; map to physical device, then set
+              env var to this single device.
+            - None/"cpu": keep default visibility.
+            - Otherwise: set env var to the provided single device string.
+        device_type: Device type ("cuda" or "npu"). If None, auto-detects.
+
+    Behavior:
+        - CUDA: Sets CUDA_VISIBLE_DEVICES and calls torch.cuda.set_device()
+        - NPU: Sets ASCEND_RT_VISIBLE_DEVICES and calls torch.npu.set_device()
     """
+    from vllm_omni.utils import detect_device_type, get_device_control_env_var
+
+    if device_type is None:
+        device_type = detect_device_type()
+
+    env_var = get_device_control_env_var()
+
+    # Select device-specific torch functions
+    if device_type == "npu":
+        try:
+            import torch.npu  # type: ignore[import-untyped]
+        except ImportError:
+            logger.debug("[Stage-%s] torch.npu not available, skipping NPU device setup", stage_id)
+            return
+
+        is_available_fn = torch.npu.is_available
+        set_device_fn = torch.npu.set_device
+        device_count_fn = torch.npu.device_count
+        get_device_properties_fn = torch.npu.get_device_properties
+        mem_get_info_fn = torch.npu.mem_get_info
+        get_device_name_fn = torch.npu.get_device_name
+        device_type_label = "NPU"
+    elif device_type == "cuda":
+        import torch  # noqa: WPS433
+
+        is_available_fn = torch.cuda.is_available
+        set_device_fn = torch.cuda.set_device
+        device_count_fn = torch.cuda.device_count
+        get_device_properties_fn = torch.cuda.get_device_properties
+        mem_get_info_fn = torch.cuda.mem_get_info
+        get_device_name_fn = torch.cuda.get_device_name
+        device_type_label = "CUDA"
+    else:
+        logger.debug("[Stage-%s] Unsupported device type: %s", stage_id, device_type)
+        return
+
     try:
         selected_physical: int | None = None
         logical_idx: int | None = None
 
         if isinstance(devices, str) and "," in devices:
-            os.environ["CUDA_VISIBLE_DEVICES"] = devices
+            os.environ[env_var] = 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-%s] Set %s to %s; logical 0 -> physical %s",
                         stage_id,
+                        env_var,
                         devices,
                         selected_physical,
                     )
                 except Exception as e:
-                    logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
+                    logger.debug("[Stage-%s] Failed to parse first %s device: %s", stage_id, device_type_label, 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")
+            vis = os.environ.get(env_var)
             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)
+                    logger.debug("[Stage-%s] Failed to map logical index via %s: %s", stage_id, env_var, e)
                     selected_physical = None
             if selected_physical is None:
                 selected_physical = int(logical_idx)
-            os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
+            os.environ[env_var] = str(selected_physical)
             logger.debug(
-                "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES to single device",
+                "[Stage-%s] Logical index %d -> physical %s; set %s to single device",
                 stage_id,
                 logical_idx + 1,
                 selected_physical,
+                env_var,
             )
         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)
+            os.environ[env_var] = str(selected_physical)
-            logger.debug(
+            logger.debug("[Stage-%s] Set %s to single device %s (fallback)", stage_id, env_var, selected_physical)
-                "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical
-            )
 
         try:
-            import torch
+            import torch  # noqa: WPS433
 
-            if torch.cuda.is_available():
+            if is_available_fn():
                 try:
-                    torch.cuda.set_device(0)
+                    set_device_fn(0)
                 except Exception as e:
                     logger.debug(
-                        "[Stage-%s] torch.cuda.set_device(0) failed: %s",
+                        "[Stage-%s] %s set_device(0) failed: %s", stage_id, device_type_label, e, exc_info=True
-                        stage_id,
-                        e,
-                        exc_info=True,
                     )
-                num = torch.cuda.device_count()
+                num = device_count_fn()
                 info = []
                 for i in range(num):
-                    total = torch.cuda.get_device_properties(i).total_memory
+                    total = get_device_properties_fn(i).total_memory
-                    free, _ = torch.cuda.mem_get_info(i)
+                    free, _ = mem_get_info_fn(i)
                     info.append(
                         {
                             "idx": i,
-                            "name": torch.cuda.get_device_name(i),
+                            "name": get_device_name_fn(i),
                             "total": int(total),
                             "free": int(free),
                         }
                     )
-                logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info)
+                logger.debug("[Stage-%s] %s devices visible=%s info=%s", stage_id, device_type_label, num, info)
         except Exception as e:
-            logger.debug(
+            logger.debug("[Stage-%s] Failed to query %s devices: %s", stage_id, device_type_label, e, exc_info=True)
-                "[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)
 

vllm_omni/entrypoints/utils.py

@@ -1,24 +1,56 @@
-from __future__ import annotations
-
-import logging
 import os
 from pathlib import Path
 
 from omegaconf import OmegaConf
 from vllm.transformers_utils.config import get_config
 
-logger = logging.getLogger(__name__)
+from vllm_omni.utils import detect_device_type
 
 # Get the project root directory (2 levels up from this file)
 PROJECT_ROOT = Path(__file__).parent.parent.parent
 
 
+def resolve_model_config_path(model: str) -> str:
+    """Resolve the stage config file path from the model name.
+
+    Resolves stage configuration path based on the model type and device type.
+    First tries to find a device-specific YAML file from stage_configs/{device_type}/
+    directory. If not found, falls back to the default config file.
+
+    Args:
+        model: Model name or path (used to determine model_type)
+
+    Returns:
+        String path to the stage configuration file
+
+    Raises:
+        FileNotFoundError: If no stage config file exists for the model type
+    """
+    hf_config = get_config(model, trust_remote_code=True)
+    model_type = hf_config.model_type
+    device_type = detect_device_type()
+
+    # Try device-specific config first
+    if device_type != "cuda":
+        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
+        device_config_path = PROJECT_ROOT / device_config_file
+        if os.path.exists(device_config_path):
+            return str(device_config_path)
+
+    # Fall back to default config
+    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
+    stage_config_path = PROJECT_ROOT / stage_config_file
+    if not os.path.exists(stage_config_path):
+        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
+    return str(stage_config_path)
+
+
 def load_stage_configs_from_model(model: str) -> list:
     """Load stage configurations from model's default config file.
 
-    Loads stage configurations based on the model type. Looks for a
+    Loads stage configurations based on the model type and device type.
-    YAML configuration file in the stage_configs directory matching
+    First tries to load a device-specific YAML file from stage_configs/{device_type}/
-    the model's model_type.
+    directory. If not found, falls back to the default config file.
 
     Args:
         model: Model name or path (used to determine model_type)
@@ -29,13 +61,8 @@ def load_stage_configs_from_model(model: str) -> list:
     Raises:
         FileNotFoundError: If no stage config file exists for the model type
     """
-    hf_config = get_config(model, trust_remote_code=True)
+    stage_config_path = resolve_model_config_path(model)
-    model_type = hf_config.model_type
+    stage_configs = load_stage_configs_from_yaml(config_path=stage_config_path)
-    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
-    stage_config_path = PROJECT_ROOT / stage_config_file
-    if not os.path.exists(stage_config_path):
-        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
-    stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path))
     return stage_configs
 
 

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -35,6 +35,7 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
 )
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
 from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision
+from vllm_omni.utils.platform_utils import is_npu
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -201,6 +202,7 @@ class Qwen2_5OmniForConditionalGeneration(
         """
         if self.model_stage == "thinker":
             # Normalize to batched inputs if caller provides 1D/2D unbatched tensors
+            # TODO: Remove this hack when NPU supports batched inputs properly
             added_batch_dim = False
             if input_ids is not None and input_ids.ndim == 1:
                 input_ids = input_ids.unsqueeze(0)
@@ -228,12 +230,25 @@ class Qwen2_5OmniForConditionalGeneration(
                 positions = positions.to(thinker_dev)
             if inputs_embeds is not None and inputs_embeds.device != thinker_dev:
                 inputs_embeds = inputs_embeds.to(thinker_dev)
+
+            if is_npu():
+                # TODO: remove this hack when NPU supports batched inputs properly
+                thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids
+                thinker_positions = positions[0] if positions.ndim > 1 else positions
+                thinker_inputs_embeds = (
+                    inputs_embeds[0] if inputs_embeds is not None and added_batch_dim else inputs_embeds
+                )
+            else:
+                thinker_input_ids = input_ids
+                thinker_positions = positions[0]
+                thinker_inputs_embeds = inputs_embeds
+
             # Run thinker
             thinker_output = self.thinker(
-                input_ids=input_ids,
+                input_ids=thinker_input_ids,
-                positions=positions[0],
+                positions=thinker_positions,
                 intermediate_tensors=intermediate_tensors,
-                inputs_embeds=inputs_embeds,
+                inputs_embeds=thinker_inputs_embeds,
                 **kwargs,
             )
 
@@ -851,7 +866,7 @@ class Qwen2_5OmniForConditionalGeneration(
         __init__."""
         if self.token2wav is None or self.token2wav_config is None:
             return
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = self._module_device(self.token2wav)
         # optional speaker resources
         conds = getattr(self.token2wav_config, "conds", None)
         ref_mels = getattr(self.token2wav_config, "ref_mels", None)

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py

@@ -33,6 +33,8 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.utils.platform_utils import is_npu
+
 
 # Provide a no-op auto_docstring decorator to satisfy annotations if missing
 def auto_docstring(func=None, **_kwargs):
@@ -724,7 +726,13 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
     else:
         beta = 0.0
 
-    kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
+    # TODO: When torch.kaiser_window supports NPU, remove the device="cpu" argument
+    if is_npu():
+        kaiser_window = torch.kaiser_window(
+            kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
+        ).to("npu")
+    else:
+        kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
 
     # Compute time indices
     if is_even:
@@ -734,7 +742,7 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
 
     # Compute sinc filter
     if cutoff == 0:
-        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)  # Ensures correct shape
+        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)
 
     sinc_filter = torch.sinc(2 * cutoff * time_indices)
     normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
@@ -745,6 +753,29 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
     return normalized_filter.view(1, 1, kernel_size)
 
 
+def replication_pad_1d(hidden_states: torch.Tensor, pad_left: int, pad_right: int) -> torch.Tensor:
+    """
+    Manual replicate padding to avoid replication_pad1d kernel limits on NPU.
+    TODO: remove when F.pad supports replicate mode on NPU.
+    """
+    # NOTE: a immature implementation for running in NPU. Need to discuss.
+    if pad_left == 0 and pad_right == 0:
+        return hidden_states
+
+    segments = []
+    if pad_left > 0:
+        left = hidden_states[..., :1].expand(*hidden_states.shape[:-1], pad_left)
+        segments.append(left)
+
+    segments.append(hidden_states)
+
+    if pad_right > 0:
+        right = hidden_states[..., -1:].expand(*hidden_states.shape[:-1], pad_right)
+        segments.append(right)
+
+    return torch.cat(segments, dim=-1)
+
+
 class UpSample1d(nn.Module):
     def __init__(self, ratio=2, kernel_size=None):
         super().__init__()
@@ -760,14 +791,28 @@ class UpSample1d(nn.Module):
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
-        hidden_states_dtype = hidden_states.dtype
+        if is_npu():
-        hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype)
+            # TODO: When F.pad supports replicate mode on NPU, remove this branch
-        hidden_states = self.ratio * F.conv_transpose1d(
+            input_dtype = hidden_states.dtype
-            hidden_states,
+            # F.pad in NPU doesn't support BF16 when mode is replicate.
-            self.filter.expand(channels, -1, -1),
+            # To ensure the accuracy, manually pad the input tensor.
-            stride=self.stride,
+            hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad, self.pad)
-            groups=channels,
+            filter_convert_dtype = self.filter.to(hidden_states.dtype)
-        ).to(hidden_states_dtype)
+            hidden_states = self.ratio * F.conv_transpose1d(
+                hidden_states,
+                filter_convert_dtype.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(input_dtype)
+        else:
+            hidden_states_dtype = hidden_states.dtype
+            hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype)
+            hidden_states = self.ratio * F.conv_transpose1d(
+                hidden_states,
+                self.filter.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(hidden_states_dtype)
         hidden_states = hidden_states[..., self.pad_left : -self.pad_right]
 
         return hidden_states
@@ -793,14 +838,29 @@ class DownSample1d(nn.Module):
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
-        hidden_states_dtype = hidden_states.dtype
+        if is_npu():
-        hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(self.filter.dtype)
+            input_dtype = hidden_states.dtype
-        out = F.conv1d(
+            # F.pad in NPU doesn't support BF16 when mode is replicate.
-            hidden_states,
+            # To ensure the accuracy, manually pad the input tensor.
-            self.filter.expand(channels, -1, -1),
+            hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad_left, self.pad_right)
-            stride=self.stride,
+            filter_on_device = self.filter.to(device=hidden_states.device, dtype=hidden_states.dtype)
-            groups=channels,
+            out = F.conv1d(
-        ).to(hidden_states_dtype)
+                hidden_states,
+                filter_on_device.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(input_dtype)
+        else:
+            hidden_states_dtype = hidden_states.dtype
+            hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(
+                self.filter.dtype
+            )
+            out = F.conv1d(
+                hidden_states,
+                self.filter.expand(channels, -1, -1),
+                stride=self.stride,
+                groups=channels,
+            ).to(hidden_states_dtype)
         return out
 
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -519,9 +519,9 @@ class Qwen3OmniMoeForConditionalGeneration(
                         tts_eos_thinker=tts_eos_thinker,
                         tts_pad_thinker=tts_pad_thinker,
                     )
-                    # req_embeds: [S,D], req_input_ids: [S]  →  map到当前 span
+                    # req_embeds: [S,D], req_input_ids: [S]  →  map onto the current span
                     seg_len = min(span_len, req_embeds.shape[0])
-                    # 强制对齐 dtype/device/内存布局，避免精度或布局差异导致比较不相等
+                    # Force-align dtype/device/memory layout to avoid mismatches from precision or layout differences
                     inputs_embeds[s : s + seg_len] = (
                         req_embeds[:seg_len].to(device=inputs_embeds.device, dtype=inputs_embeds.dtype).contiguous()
                     )

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py

@@ -6,6 +6,7 @@ The code predictor generates residual RVQ (Residual Vector Quantization) codes
 autoregressively, predicting layers 1 to N based on layer-0 codes from the talker.
 """
 
+from collections import namedtuple
 from typing import Any, Optional
 
 import torch
@@ -164,21 +165,48 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
 
         from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 
-        attention_interface = ALL_ATTENTION_FUNCTIONS["flash_attention_2"]
+        # Try attention backends in order of preference, with runtime error handling
-        attn_output, _ = attention_interface(
+        # This handles cases where the backend is registered but not actually available
-            self,
+        attention_backends = ["flash_attention_2", "xformers", "eager", "sdpa"]
-            q_heads,
+        attn_output = None
-            k_heads,
+        last_error = None
-            v_heads,
+
-            None,
+        for backend_name in attention_backends:
-            dropout=0.0 if not self.training else self.attention_dropout,
+            if backend_name not in ALL_ATTENTION_FUNCTIONS:
-            scaling=self.head_dim**-0.5,
+                continue
-            sliding_window=None,
+
-            use_cache=use_cache,
+            try:
-            position_ids=position_ids,
+                attention_interface = ALL_ATTENTION_FUNCTIONS[backend_name]
-            output_hidden_states=True,
+                attn_output, _ = attention_interface(
-            output_attentions=False,
+                    self,
-        )
+                    q_heads,
+                    k_heads,
+                    v_heads,
+                    None,
+                    dropout=0.0 if not self.training else getattr(self, "attention_dropout", 0.0),
+                    scaling=self.head_dim**-0.5,
+                    sliding_window=None,
+                    use_cache=use_cache,
+                    position_ids=position_ids,
+                    output_hidden_states=True,
+                    output_attentions=False,
+                )
+                # Success - log fallback if not using flash_attention_2
+                if backend_name != "flash_attention_2":
+                    logger.warning_once(
+                        f"Using {backend_name} attention backend (flash_attention_2 not available or failed)"
+                    )
+                break
+            except (ValueError, ImportError, RuntimeError, AttributeError) as e:
+                # Store error and try next backend
+                last_error = e
+                continue
+
+        if attn_output is None:
+            raise RuntimeError(
+                f"All attention backends failed. Last error: {last_error}. "
+                "Please install flash-attn, or ensure PyTorch's scaled_dot_product_attention is available."
+            )
         attn_output = attn_output.reshape(*(hidden_states.shape[:-1]), -1).contiguous()
 
         attn_output = self.o_proj(attn_output)
@@ -373,16 +401,22 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
         past_key_values: Optional[Any] = None,
         use_cache: Optional[bool] = False,
         cache_position: Optional[torch.LongTensor] = None,
-        **kwargs,
+        **kwargs: Any,
-    ):
+    ) -> Any:
         """
         Forward pass matching HF structure.
 
         Args:
             inputs_embeds: [batch, seq_len, hidden_size]
+            attention_mask: Optional attention mask tensor
+            position_ids: Optional position IDs tensor
+            past_key_values: Optional cached key-value pairs
+            use_cache: Whether to use cache
+            cache_position: Optional cache position tensor
+            **kwargs: Additional keyword arguments
 
         Returns:
-            Object with .last_hidden_state attribute
+            Named tuple with .last_hidden_state and .past_key_values attributes
         """
         batch_size, seq_len, _ = inputs_embeds.shape
 
@@ -421,8 +455,6 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
         hidden_states = self.norm(hidden_states)
 
         # Return in HF-compatible format
-        from collections import namedtuple
-
         Output = namedtuple("Output", ["last_hidden_state", "past_key_values"])
         return Output(last_hidden_state=hidden_states, past_key_values=None)  # [batch, num_code_groups-1, hidden_size]
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -1,4 +1,5 @@
 from collections.abc import Iterable
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -284,7 +285,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
 
         return result_codes, summed_embeddings
 
-    def init_multi_modal(self, thinker_config):
+    def init_multi_modal(self, thinker_config: Any) -> None:
         """
         Initialize multimodal components from the thinker.
 

vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml

@@ -0,0 +1,94 @@
+# stage config for running qwen2.5-omni with architecture of OmniLLM.
+stage_args:
+  - stage_id: 0
+    runtime:
+      process: true            # Run this stage in a separate process
+      devices: "0"            # Visible devices for this stage
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true  # Now we only support eager mode
+      trust_remote_code: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+  - stage_id: 1
+    runtime:
+      process: true
+      devices: "1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: talker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: latent
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
+    default_sampling_params:
+      temperature: 0.9
+      top_p: 0.8
+      top_k: 40
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+      stop_token_ids: [8294]
+  - stage_id: 2
+    runtime:
+      process: true
+      devices: "0"            # Example: use a different NPU than the previous stage; use "0" if single NPU
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.npu.npu_diffusion_worker.NPUDiffusionWorker
+      scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
+      gpu_memory_utilization: 0.15
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio
+    engine_input_source: [1]
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1             # Simplified: trigger downstream only after full upstream completion
+    max_inflight: 1             # Simplified: process serially within each stage
+  edges:
+    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
+      to: 1
+      window_size: -1
+    - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
+      to: 2
+      window_size: -1

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -26,6 +26,9 @@ stage_args:
       seed: 42
       detokenize: True
       repetition_penalty: 1.1
+    # Distributed connector configuration (optional)
+    output_connectors:
+      to_stage_1: connector_thinker2talker # thinker -> talker: use local mooncake
   - stage_id: 1
     runtime:
       process: true
@@ -52,17 +55,22 @@ stage_args:
       detokenize: True
       repetition_penalty: 1.05
       stop_token_ids: [8294]
+    # Distributed connector configuration (optional)
+    input_connectors:
+      from_stage_0: connector_thinker2talker # talker <- thinker: use local mooncake
+    output_connectors:
+      to_stage_2: connector_talker2generator # talker -> code2wav: use remote mooncake (example)
   - stage_id: 2
     runtime:
       process: true
-      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      devices: "2"            # Example: use a different GPU than the previous stage; use "0" if single GPU
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
-      gpu_memory_utilization: 0.15
+      gpu_memory_utilization: 0.3
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false
@@ -78,6 +86,9 @@ stage_args:
       seed: 42
       detokenize: True
       repetition_penalty: 1.1
+    # Distributed connector configuration (optional)
+    input_connectors:
+      from_stage_1: connector_talker2generator # code2wav <- talker: use remote mooncake
 
 # Top-level runtime config (concise): default windows and stage edges
 runtime:
@@ -85,6 +96,31 @@ runtime:
   defaults:
     window_size: -1             # Simplified: trigger downstream only after full upstream completion
     max_inflight: 1             # Simplified: process serially within each stage
+
+  # Distributed connectors configuration (optional)
+  connectors:
+    # Local mooncake for same-node communication
+    connector_thinker2talker:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://10.90.67.86:8080/metadata"
+        master: "10.90.67.86:50051"
+        segment: 512000000    # 512MB
+        localbuf: 64000000     # 64MB
+        proto: "tcp"
+
+    # Remote mooncake for cross-node communication (example)
+    connector_talker2generator:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://10.90.67.86:8080/metadata"
+        master: "10.90.67.86:50051"
+        segment: 1073741824    # 1GB
+        localbuf: 134217728     # 128MB
+        proto: "tcp"
+
   edges:
     - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
       to: 1

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -3,7 +3,7 @@
 # Copyright 2025 The Qwen team.
 """Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition."""
 
-from typing import Union
+from typing import Any, Union
 
 import torch
 from vllm.inputs import TextPrompt
@@ -47,11 +47,11 @@ def _compute_talker_prompt_ids_length(info):
 
 
 def thinker2talker(
-    stage_list,
+    stage_list: list[Any],
-    engine_input_source,
+    engine_input_source: list[int],
-    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
+    prompt: Union[OmniTokensPrompt, TextPrompt, None] = None,
     requires_multimodal_data: bool = False,
-):
+) -> list[OmniTokensPrompt]:
     """
     Process thinker outputs to create talker inputs.
 
@@ -64,6 +64,7 @@ def thinker2talker(
         stage_list: List of stage objects
         engine_input_source: Source stage IDs (typically [0] for thinker)
         prompt: Original prompt data
+        requires_multimodal_data: Whether multimodal data is required
 
     Returns:
         List of OmniTokensPrompt for talker stage
@@ -111,11 +112,11 @@ def thinker2talker(
 
 
 def talker2code2wav(
-    stage_list,
+    stage_list: list[Any],
-    engine_input_source,
+    engine_input_source: list[int],
-    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
+    prompt: Union[OmniTokensPrompt, TextPrompt, None] = None,
     requires_multimodal_data: bool = False,
-):
+) -> list[OmniTokensPrompt]:
     """
     Process talker outputs to create code2wav inputs.
 
@@ -128,6 +129,7 @@ def talker2code2wav(
         stage_list: List of stage objects
         engine_input_source: Source stage IDs (typically [1] for talker)
         prompt: Original prompt data
+        requires_multimodal_data: Whether multimodal data is required
 
     Returns:
         List of OmniTokensPrompt for code2wav stage

vllm_omni/utils/__init__.py

@@ -0,0 +1,11 @@
+from vllm_omni.utils.platform_utils import (
+    detect_device_type,
+    get_device_control_env_var,
+    is_npu,
+)
+
+__all__ = [
+    "detect_device_type",
+    "get_device_control_env_var",
+    "is_npu",
+]

vllm_omni/utils/platform_utils.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import torch
+from vllm.platforms import current_platform
+
+
+def detect_device_type() -> str:
+    device_type = getattr(current_platform, "device_type", None)
+    if isinstance(device_type, str) and device_type:
+        return device_type.lower()
+    if torch.cuda.is_available():
+        return "cuda"
+    if hasattr(torch, "npu") and torch.npu.is_available():  # type: ignore[attr-defined]
+        return "npu"
+    return "cpu"
+
+
+def is_npu() -> bool:
+    return detect_device_type() == "npu"
+
+
+def get_device_control_env_var() -> str:
+    """Return the environment variable name for device visibility control."""
+    if hasattr(current_platform, "device_control_env_var"):
+        env_var = getattr(current_platform, "device_control_env_var", None)
+        if isinstance(env_var, str) and env_var:
+            return env_var
+
+    device_type = detect_device_type()
+    if device_type == "npu":
+        return "ASCEND_RT_VISIBLE_DEVICES"
+    return "CUDA_VISIBLE_DEVICES"  # fallback
+
+
+def get_diffusion_worker_class() -> type:
+    """Get the appropriate diffusion WorkerProc class based on current device type.
+
+    Returns:
+        The WorkerProc class for the detected device type (either NPUWorkerProc or WorkerProc).
+
+    Raises:
+        ImportError: If the worker module for the detected device is not available.
+    """
+    device_type = detect_device_type()
+
+    if device_type == "npu":
+        from vllm_omni.diffusion.worker.npu.npu_worker import NPUWorkerProc
+
+        return NPUWorkerProc
+    else:
+        # Default to GPU worker for cuda and other devices
+        from vllm_omni.diffusion.worker.gpu_worker import WorkerProc
+
+        return WorkerProc

vllm_omni/worker/gpu_generation_model_runner.py

@@ -1,4 +1,4 @@
-"""Code2Wav GPU Model Runner for vLLM-omni.
+"""Code2Wav GPU Model Runner for vLLM-Omni.
 
 Handles direct conversion from codec codes to audio waveforms for Qwen3 Omni MoE Code2Wav.
 This is a non-autoregressive model that doesn't require sampling or logits computation.
@@ -37,7 +37,7 @@ logger = logging.getLogger(__name__)
 
 
 class GPUGenerationModelRunner(OmniGPUModelRunner):
-    """Generation model runner for vLLM-omni (non-autoregressive).
+    """Generation model runner for vLLM-Omni (non-autoregressive).
 
     - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
     - Does not compute logits or perform token sampling.

vllm_omni/worker/gpu_model_runner.py

@@ -6,6 +6,8 @@ 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.layers.rotary_embedding import MRotaryEmbedding
+from vllm.model_executor.models.interfaces import supports_mrope
 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
 from vllm.sampling_params import SamplingType
 from vllm.utils import LazyLoader, cdiv
@@ -33,6 +35,63 @@ logger = init_logger(__name__)
 
 
 class OmniGPUModelRunner(GPUModelRunner):
+    def _init_mrope_positions(self, req_state: CachedRequestState):
+        """Initialize M-RoPE positions for multimodal inputs.
+
+        Extracts multimodal feature metadata (image grids, video grids,
+        audio features) and computes M-RoPE positions for proper positional
+        encoding of multimodal tokens.
+
+        Args:
+            req_state: Cached request state containing multimodal features
+
+        Raises:
+            AssertionError: If the model does not support M-RoPE
+        """
+        image_grid_thw = []
+        video_grid_thw = []
+        second_per_grid_ts = []
+        audio_feature_lengths = []
+        use_audio_in_video = False
+        for mm_feature in req_state.mm_features:
+            mm_item = mm_feature.data
+            if mm_item is None:
+                continue
+            mm_input = mm_item.get_data()
+            if (t := mm_input.get("image_grid_thw")) is not None:
+                image_grid_thw.append(t.tolist())
+            if (t := mm_input.get("video_grid_thw")) is not None:
+                video_grid_thw.append(t.tolist())
+            if (t := mm_input.get("second_per_grid_ts")) is not None:
+                second_per_grid_ts.append(t)
+            if (t := mm_input.get("audio_feature_lengths")) is not None:
+                audio_feature_lengths.append(t)
+            # Check for use_audio_in_video
+            use_audio_in_video_value = mm_input.get("use_audio_in_video")
+            if use_audio_in_video_value is not None:
+                use_audio_in_video = bool(use_audio_in_video_value.item())
+
+        if supports_mrope(self.model):
+            req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions(
+                req_state.prompt_token_ids,
+                hf_config=self.model_config.hf_config,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                second_per_grid_ts=second_per_grid_ts,
+                audio_feature_lengths=audio_feature_lengths,
+                use_audio_in_video=use_audio_in_video,
+            )
+        else:
+            req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor(
+                req_state.prompt_token_ids,
+                hf_config=self.model_config.hf_config,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                second_per_grid_ts=second_per_grid_ts,
+                audio_feature_lengths=audio_feature_lengths,
+                use_audio_in_video=use_audio_in_video,
+            )
+
     def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
         """Update the cached states and the persistent batch with the scheduler
         output.

vllm_omni/worker/npu/npu_ar_model_runner.py

@@ -0,0 +1,963 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from vllm.config import CUDAGraphMode
+from vllm.distributed import tensor_model_parallel_all_gather
+from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
+from vllm.distributed.parallel_state import get_pp_group, get_tp_group
+from vllm.forward_context import BatchDescriptor, get_forward_context
+from vllm.logger import logger
+from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
+from vllm.v1.core.sched.output import SchedulerOutput
+
+# yapf conflicts with isort for this block
+# yapf: disable
+from vllm.v1.kv_cache_interface import (
+    EncoderOnlyAttentionSpec,
+)
+
+# yapf: enable
+from vllm.v1.outputs import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    AsyncModelRunnerOutput,
+    ModelRunnerOutput,
+)
+from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
+from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput
+from vllm_ascend.ascend_forward_context import set_ascend_forward_context
+from vllm_ascend.attention.attention_v1 import AscendAttentionState
+from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
+from vllm_ascend.compilation.acl_graph import update_attn_params, update_mla_attn_params
+from vllm_ascend.spec_decode.interface import SpecDcodeType
+from vllm_ascend.utils import (
+    ProfileExecuteDuration,
+    enable_sp,
+    lmhead_tp_enable,
+)
+from vllm_ascend.worker.model_runner_v1 import AsyncNPUModelRunnerOutput
+
+from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner
+
+
+class NPUARModelRunner(OmniNPUModelRunner):
+    """Autoregressive NPU model runner that returns hidden states per request."""
+
+    def _prepare_inputs(
+        self,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+    ) -> tuple[
+        dict[str, Any],
+        torch.Tensor,
+        np.ndarray,
+        int,
+        torch.Tensor,
+        int,
+        torch.Tensor,
+        SpecDecodeMetadata,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        int,
+        dict[str, dict] | None,
+    ]:
+        total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        assert total_num_scheduled_tokens > 0
+        num_reqs = self.input_batch.num_reqs
+        assert num_reqs > 0
+
+        # OPTIMIZATION: Start copying the block table first.
+        # This way, we can overlap the copy with the following CPU operations.
+        self.input_batch.block_table.commit_block_table(num_reqs)
+
+        # Get the number of scheduled tokens for each request.
+        req_ids = self.input_batch.req_ids
+        tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
+        num_scheduled_tokens = np.array(tokens, dtype=np.int32)
+        max_num_scheduled_tokens = num_scheduled_tokens.max()
+        num_valid_tokens = np.array(
+            [
+                num_tokens - len(scheduler_output.scheduled_spec_decode_tokens.get(i, []))
+                for num_tokens, i in zip(tokens, req_ids)
+            ],
+            dtype=np.int32,
+        )
+
+        if self.use_aclgraph and total_num_scheduled_tokens <= self.aclgraph_batch_sizes[-1]:
+            # Add padding to the batch size.
+            num_input_tokens = self.vllm_config.pad_for_cudagraph(total_num_scheduled_tokens)
+        elif self.use_aclgraph and enable_sp(self.vllm_config):
+            # When using aclgraph, if total_num_scheduled_tokens exceeds the maximum graph size,
+            # the model will fall back to running its FX graph in eager mode.
+            # In this case, when sequence parallelism is enabled, we need to pad tokens to align
+            # with tp_size because pad_size cannot be captured by the FX graph
+            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
+            num_input_tokens = math.ceil(total_num_scheduled_tokens / tp_size) * tp_size
+        else:
+            # Eager mode.
+            num_input_tokens = total_num_scheduled_tokens
+
+        # Get the attention state.
+        attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, num_valid_tokens)
+        self.attn_state = attn_state  # type: ignore
+
+        # Determine if it's a splitfuse batch
+        with_prefill = attn_state not in [AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding]
+
+        self.query_lens = torch.from_numpy(num_scheduled_tokens)
+        enable_dbo = self._check_dbo_is_valid(self.query_lens.tolist(), attn_state, total_num_scheduled_tokens)
+
+        # Get info across DP ranks.
+        # NOTE: maybe_padded_num_tokens is only used when using TorchAir with DP,
+        # Otherwise, it's just max_tokens_across_dp_cpu
+        (maybe_padded_num_tokens, num_tokens_across_dp, with_prefill, enable_dbo) = self._sync_metadata_across_dp(
+            num_input_tokens, with_prefill, enable_dbo
+        )
+
+        # TODO: Now that num_input_tokens is basically identical with maybe_padded_num_tokens
+        # We should consider removing maybe_padded_num_tokens later
+        num_input_tokens = maybe_padded_num_tokens
+
+        # Hot-Swap lora model
+        if self.lora_config:
+            self.set_active_loras(self.input_batch, num_scheduled_tokens)
+
+        # Get request indices.
+        # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
+        req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)
+
+        # cu_num_tokens: [2, 5, 3] -> [2, 7, 10]
+        # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
+        cu_num_tokens, arange = self._get_cumsum_and_arange(num_scheduled_tokens)
+
+        positions_np = self.positions_np[:total_num_scheduled_tokens]
+        np.add(self.input_batch.num_computed_tokens_cpu[req_indices], arange, out=positions_np)
+
+        # Calculate M-RoPE positions.
+        # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
+        if self.uses_mrope:
+            self._calc_mrope_positions(scheduler_output)
+
+            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
+            self.mrope_positions[:, :total_num_scheduled_tokens].copy_(
+                self.mrope_positions_cpu[:, :total_num_scheduled_tokens], non_blocking=True
+            )
+
+        # Get token indices.
+        # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
+        # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]
+        # where M is the max_model_len.
+        token_indices = positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]
+
+        # Prepare input_ids.
+        # NOTE(woosuk): We use torch.index_select instead of np.take here
+        # because torch.index_select is much faster than np.take for large
+        # tensors.
+        torch.index_select(
+            self.input_batch.token_ids_cpu_tensor.flatten(),
+            0,
+            torch.from_numpy(token_indices),
+            out=self.input_ids_cpu[:total_num_scheduled_tokens],
+        )
+
+        # Prepare some information for building Attention-Metadata
+        # Compute and commit slot mapping
+        self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np)
+        self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens)
+
+        self.query_start_loc_np[0] = 0
+        self.query_start_loc_np[1 : num_reqs + 1] = cu_num_tokens
+        self.query_start_loc[: num_reqs + 1].copy_(self.query_start_loc_cpu[: num_reqs + 1], non_blocking=True)
+
+        self.seq_lens_np[:num_reqs] = self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens
+        self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs], non_blocking=True)
+
+        # Fill unused with -1. Needed for reshape_and_cache
+        self.query_start_loc[num_reqs + 1 :].fill_(-1)
+        self.seq_lens[num_reqs:].fill_(0)
+
+        self.query_lens = torch.from_numpy(num_scheduled_tokens)
+
+        # Copy the tensors to the NPU.
+        self._prepare_input_ids(total_num_scheduled_tokens, cu_num_tokens)
+        self.positions_cpu[total_num_scheduled_tokens:num_input_tokens].zero_()
+        self.positions[:num_input_tokens].copy_(self.positions_cpu[:num_input_tokens], non_blocking=True)
+
+        # Make Attention metadata
+        positions_cpu = self.positions_cpu[:num_input_tokens]
+        positions = self.positions[:num_input_tokens]
+        seq_lens_cpu = self.seq_lens_cpu[:num_reqs]
+        attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, num_valid_tokens)
+        self.attn_mask = self._make_attention_mask(seq_lens=seq_lens_cpu, position=positions_cpu, attn_state=attn_state)
+        self.attn_state = attn_state  # type: ignore
+
+        self.with_prefill = with_prefill
+        self.num_tokens_across_dp = num_tokens_across_dp
+        self._update_graph_pad_size(with_prefill, maybe_padded_num_tokens)
+        attn_metadata: dict[str, Any] = {}
+
+        # Omni-new: per_req_additional_information
+        per_req_additional_information: dict[str, dict] | None = None
+
+        # _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)
+
+            # 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.
+            input_ids = self.input_ids[:total_num_scheduled_tokens]
+            if mm_embeds:
+                inputs_embeds = self.model.get_input_embeddings(input_ids, mm_embeds)
+            else:
+                inputs_embeds = self.model.get_input_embeddings(input_ids)
+            # TODO(woosuk): Avoid the copy. Optimize.
+            self.inputs_embeds[:total_num_scheduled_tokens].copy_(inputs_embeds)
+
+            #  -------------------------------------- Omni-new -------------------------------------------------
+            # Omni-new: Reset per-step additional information collector (deprecated concat path)
+            if hasattr(self, "_forward_additional_information"):
+                self._forward_additional_information = None
+            # Omni-new: per-request additional information for this step
+            per_req_additional_information = {}
+
+            # Omni-new: 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[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
+            #  ------------------------------------------------------------------------------------------------
+
+            inputs_embeds = self.inputs_embeds[:num_input_tokens]
+            input_ids = self.input_ids[: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 ACL graph.
+            input_ids = self.input_ids[:num_input_tokens]
+            inputs_embeds = None
+        positions = self.positions[:num_input_tokens]
+        input_ids, positions = self._update_input_ids_and_positions(
+            input_ids, positions, num_input_tokens, with_prefill, maybe_padded_num_tokens
+        )
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            assert intermediate_tensors is not None
+            assert self.intermediate_tensors is not None
+            for k, v in intermediate_tensors.items():
+                self.intermediate_tensors[k][:num_input_tokens].copy_(v[:num_input_tokens], non_blocking=True)
+            intermediate_tensors = IntermediateTensors(
+                {k: v[:num_input_tokens] for k, v in self.intermediate_tensors.items()}
+            )
+
+        use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
+        if not use_spec_decode:
+            # NOTE(woosuk): Due to chunked prefills, the batch may contain
+            # partial requests. While we should not sample any token
+            # from these partial requests, we do so for simplicity.
+            # We will ignore the sampled tokens from the partial requests.
+            # TODO: Support prompt logprobs.
+            spec_decode_metadata = None
+            logits_indices = torch.from_numpy(cu_num_tokens - 1).to(self.device, non_blocking=True)
+        else:
+            # Get the number of draft tokens for each request.
+            # Iterate over the dictionary rather than all requests since not all
+            # requests have draft tokens.
+            num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
+            for req_id, draft_token_ids in scheduler_output.scheduled_spec_decode_tokens.items():
+                req_idx = self.input_batch.req_id_to_index[req_id]
+                num_draft_tokens[req_idx] = len(draft_token_ids)
+
+            spec_decode_metadata = self._calc_spec_decode_metadata(num_draft_tokens, cu_num_tokens)
+            logits_indices = spec_decode_metadata.logits_indices
+            self.num_draft_tokens.np[:num_reqs] = num_draft_tokens
+            self.num_draft_tokens.np[num_reqs:].fill(0)
+            self.num_draft_tokens.copy_to_gpu()
+
+        # Used in the below loop.
+        # query_start_loc_cpu = self.query_start_loc.cpu[:num_reqs + 1]
+        num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs]
+        spec_decode_common_attn_metadata = None
+        if use_spec_decode and self.need_accepted_tokens:
+            self.num_accepted_tokens.np[:num_reqs] = self.input_batch.num_accepted_tokens_cpu[:num_reqs]
+            self.num_accepted_tokens.np[num_reqs:].fill(1)
+            self.num_accepted_tokens.copy_to_gpu()
+
+        # Prepare the attention metadata for each KV cache group and make layers
+        # in the same group share the same metadata.
+        for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups):
+            if isinstance(kv_cache_group_spec.kv_cache_spec, EncoderOnlyAttentionSpec):
+                # Encoder-only layers do not have KV cache, so we need to
+                # create a dummy block table and slot mapping for them.
+                blk_table_tensor = torch.zeros(
+                    (num_reqs, 1),
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+                slot_mapping = torch.zeros(
+                    (total_num_scheduled_tokens,),
+                    dtype=torch.int64,
+                    device=self.device,
+                )
+            else:
+                blk_table = self.input_batch.block_table[kv_cache_group_id]
+                blk_table_tensor = blk_table.get_device_tensor()
+                slot_mapping = blk_table.slot_mapping_cpu[:total_num_scheduled_tokens]
+                self.slot_mapping[:total_num_scheduled_tokens].copy_(
+                    slot_mapping[:total_num_scheduled_tokens],
+                    non_blocking=True,
+                )
+                self.slot_mapping[total_num_scheduled_tokens:].fill_(0)
+
+            # Make AscendCommonAttentionMetadata
+            common_attn_metadata = AscendCommonAttentionMetadata(
+                query_start_loc=self.query_start_loc_cpu[: num_reqs + 1],
+                query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1],
+                seq_lens_cpu=self.seq_lens_cpu,
+                seq_lens=self.seq_lens_cpu[:num_reqs],
+                num_reqs=num_reqs,
+                num_actual_tokens=total_num_scheduled_tokens,
+                num_input_tokens=num_input_tokens,
+                actual_seq_lengths_q=self.actual_seq_lengths_q,
+                # TODO: change this to the right block table for linear attn
+                block_table_tensor=blk_table_tensor[:num_reqs],
+                slot_mapping=self.slot_mapping,
+                num_computed_tokens_cpu=num_computed_tokens_cpu,
+                positions=self.positions,
+                attn_mask=self.attn_mask,
+                spec_attn_mask=self.spec_attn_mask,
+                attn_state=self.attn_state,
+                enable_dbo_across_dp=enable_dbo,
+                is_only_prefill=bool(np.all(num_valid_tokens != 1)),
+                max_query_len=max_num_scheduled_tokens,
+                graph_pad_size=self.graph_pad_size,
+                decode_token_per_req=self.decode_token_per_req,
+                cos=self.cos,
+                sin=self.sin,
+            )
+
+            if self.speculative_config and spec_decode_common_attn_metadata is None:
+                spec_decode_common_attn_metadata = common_attn_metadata
+
+            for attn_group in self.attn_groups[kv_cache_group_id]:
+                common_prefix_len = 0
+                extra_attn_metadata_args = {}
+                builder = attn_group.get_metadata_builder()
+                if isinstance(builder, GDNAttentionMetadataBuilder) or self.model_config.runner_type == "pooling":
+                    if use_spec_decode:
+                        extra_attn_metadata_args = dict(
+                            num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs],
+                            num_draft_tokens=self.num_draft_tokens.gpu[:num_reqs],
+                        )
+                    attn_metadata_i = builder.build(
+                        common_prefix_len=common_prefix_len,
+                        common_attn_metadata=common_attn_metadata,
+                        **extra_attn_metadata_args,
+                    )
+                else:
+                    attn_metadata_i = builder.build(
+                        common_prefix_len=common_prefix_len,
+                        common_attn_metadata=common_attn_metadata,
+                        model=self.get_model(),
+                        **extra_attn_metadata_args,
+                    )
+
+                for layer_name in attn_group.layer_names:
+                    attn_metadata[layer_name] = attn_metadata_i
+
+        if lmhead_tp_enable():
+            max_num_reqs_across_dp = maybe_padded_num_tokens if not with_prefill else self.max_num_reqs
+            logits_indices = nn.functional.pad(logits_indices, (0, max_num_reqs_across_dp - logits_indices.shape[0]))
+
+        return (
+            attn_metadata,
+            positions,
+            num_scheduled_tokens,
+            num_input_tokens,
+            num_tokens_across_dp,
+            maybe_padded_num_tokens,
+            logits_indices,
+            spec_decode_metadata,
+            input_ids,
+            inputs_embeds,
+            intermediate_tensors,
+            max_num_scheduled_tokens,
+            per_req_additional_information,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+    ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
+        with ProfileExecuteDuration().capture_async("prepare input"):
+            self._update_states(scheduler_output)
+
+            #  -------------------------------------- Omni-new -------------------------------------------------
+            # Omni-new: 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():
+                    logger.debug("skip this step for we receive the data from remote disaggregate prefill node")
+                    # Return empty ModelRunnerOutput if there's no work to do.
+                    return EMPTY_MODEL_RUNNER_OUTPUT
+                return self.kv_connector_no_forward(scheduler_output)
+
+            if self.dynamic_eplb:
+                self.eplb_updator.forward_before()
+
+            (
+                attn_metadata,
+                positions,
+                num_scheduled_tokens_np,
+                num_input_tokens,
+                num_tokens_across_dp,
+                maybe_padded_num_tokens,
+                logits_indices,
+                spec_decode_metadata,
+                input_ids,
+                inputs_embeds,
+                intermediate_tensors,
+                max_query_len,
+                per_req_additional_information,
+            ) = self._prepare_inputs(scheduler_output, intermediate_tensors)
+
+            if self.dynamic_eplb:
+                self.eplb_updator.take_update_info_from_eplb_process()
+
+        moe_comm_type = self._select_moe_comm_method(num_input_tokens, self.with_prefill)
+
+        uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
+            scheduler_output.total_num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
+        )
+        batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens, uniform_decode=uniform_decode)
+        aclgraph_runtime_mode, batch_descriptor = self.aclgraph_dispatcher.dispatch(batch_descriptor)
+
+        # Run forward pass
+        with ProfileExecuteDuration().capture_async("forward"):
+            with set_ascend_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                with_prefill=self.with_prefill,
+                reserved_mc2_mask=self.reserved_mc2_mask,
+                moe_comm_type=moe_comm_type,
+                aclgraph_runtime_mode=aclgraph_runtime_mode,
+                batch_descriptor=batch_descriptor,
+                num_actual_tokens=scheduler_output.total_num_scheduled_tokens,
+                prefetch_stream=self.prefetch_stream,
+                model_instance=self.model,
+                weight_prefetch_method=self.weight_prefetch_method,
+            ):
+                self.maybe_setup_kv_connector(scheduler_output)
+                #  -------------------------------------- Omni-new -------------------------------------------------
+                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
+                #  ------------------------------------------------------------------------------------------------
+
+                hidden_states = self._generate_process_reqs_hidden_states(
+                    attn_metadata,
+                    self.with_prefill,
+                    maybe_padded_num_tokens,
+                    input_ids,
+                    positions,
+                    intermediate_tensors,
+                    inputs_embeds,
+                    model_kwargs_extra,
+                )
+
+            self.maybe_wait_for_kv_save()
+            finished_sending, finished_recving = self.get_finished_kv_transfer(scheduler_output)
+
+            aux_hidden_states = None
+            if self.drafter and self.drafter.name == SpecDcodeType.EAGLE3:
+                hidden_states, aux_hidden_states = hidden_states
+
+        kv_connector_output = KVConnectorOutput(finished_sending=finished_sending, finished_recving=finished_recving)
+        finished_sending = None
+        finished_recving = None
+        with ProfileExecuteDuration().capture_async("post process"):
+            # 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
+
+            #  -------------------------------------- Omni-new -------------------------------------------------
+            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}"
+                )
+            #  ------------------------------------------------------------------------------------------------
+            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.
+                if not broadcast_pp_output:
+                    hidden_states.kv_connector_output = kv_connector_output
+                    return hidden_states
+                assert isinstance(hidden_states, IntermediateTensors)
+                get_pp_group().send_tensor_dict(hidden_states.tensors, all_gather_group=get_tp_group())
+                logits = None
+            else:
+                if self.input_batch.pooling_params:
+                    return self._pool(
+                        hidden_states,
+                        scheduler_output.total_num_scheduled_tokens,
+                        num_scheduled_tokens_np,
+                        finished_sending,
+                        finished_recving,
+                        kv_connector_output,
+                    )
+                sample_hidden_states = hidden_states[logits_indices]
+                logits = self.model.compute_logits(sample_hidden_states)
+            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:
+                logits = 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:
+                if lmhead_tp_enable() and logits is not None:
+                    logits = logits[: self.input_batch.num_reqs]
+                sampler_output = self.sampler(
+                    logits=logits,
+                    sampling_metadata=sampling_metadata,
+                )
+            else:
+                if lmhead_tp_enable() and logits is not None:
+                    logits = logits[: len(spec_decode_metadata.logits_indices)]
+                # 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
+                if self.need_accepted_tokens:
+                    self._update_states_after_model_execute(output_token_ids)
+
+            discard_sampled_tokens_req_indices: list[int] = []
+            # 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.
+                    # Rewind the generator state as if the token was not sampled.
+                    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)
+
+            # Copy some objects so they don't get modified after returning.
+            # This is important when using async scheduling.
+            req_ids_output_copy = self.input_batch.req_ids.copy()
+            req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
+
+            # NOTE: NPU -> 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(
+                hidden_states[: scheduler_output.total_num_scheduled_tokens],
+                scheduler_output,
+            )
+
+            num_sampled_tokens = sampler_output.sampled_token_ids.shape[0]
+            sampled_token_ids = sampler_output.sampled_token_ids
+            if not self.use_async_scheduling:
+                # Get the valid generated tokens.
+                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()
+            else:
+                valid_sampled_token_ids = []
+                invalid_req_indices = list(discard_sampled_tokens_req_indices)
+                invalid_req_indices_set = set(invalid_req_indices)
+                assert sampled_token_ids.shape[-1] == 1
+
+                # Cache the sampled tokens on the NPU and avoid CPU sync.
+                # These will be copied into input_ids in the next step
+                # when preparing inputs.
+                self.input_batch.prev_sampled_token_ids = sampled_token_ids
+                self.input_batch.prev_sampled_token_ids_invalid_indices = invalid_req_indices_set
+                self.input_batch.prev_req_id_to_index = {
+                    req_id: i for i, req_id in enumerate(self.input_batch.req_ids) if i not in invalid_req_indices_set
+                }
+            # 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 in range(num_sampled_tokens):
+                if self.use_async_scheduling:
+                    sampled_ids = [-1] * 1 if req_idx not in invalid_req_indices_set else None
+                else:
+                    sampled_ids = valid_sampled_token_ids[req_idx]
+                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.model_config.max_model_len, (
+                    "Sampled token IDs exceed the max model length. "
+                    f"Total number of tokens: {end_idx} > max_model_len: "
+                    f"{self.model_config.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 self.speculative_config:
+                self._draft_token_ids = self.propose_draft_token_ids(
+                    valid_sampled_token_ids,
+                    sampling_metadata,
+                    scheduler_output,
+                    spec_decode_metadata,
+                    positions,
+                    scheduler_output.total_num_scheduled_tokens,
+                    hidden_states,
+                    attn_metadata,
+                    aux_hidden_states,
+                )
+
+            if has_kv_transfer_group():
+                get_kv_transfer_group().clear_connector_metadata()
+
+        #  -------------------------------------- Omni-new -------------------------------------------------
+        # Omni-new: Convert to per-request tensors on CPU
+        hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
+        pooler_output: list[torch.Tensor | None] = []
+        prev_logits_index = 0
+        for rid, logits_index in zip(req_ids_output_copy, logits_indices):
+            # Base payload: hidden slice for this request in this iteration
+            hidden_slice = hidden_states_cpu[prev_logits_index : logits_index + 1]
+            payload: dict[str, object] = {"hidden": hidden_slice}
+            # Merge multimodal_outputs if present
+            if isinstance(multimodal_outputs, dict) and multimodal_outputs:
+                mm_payload: dict[str, object] = {}
+                for k, v in multimodal_outputs.items():
+                    try:
+                        # Case 1: tensor aligned on token dimension
+                        if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
+                            mm_payload[k] = v.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
+                        # Case 2: nested dict of tensors aligned on token dimension (e.g., selected_hidden_layers)
+                        elif isinstance(v, dict):
+                            sub_dict: dict[str, torch.Tensor] = {}
+                            for sk, sv in v.items():
+                                if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
+                                    sub_dict[str(sk)] = (
+                                        sv.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
+                                    )
+                            if sub_dict:
+                                mm_payload[k] = sub_dict
+                        elif isinstance(v, list):
+                            element: torch.Tensor = v[0]
+                            multimodal_outputs[k] = v[1:] if len(v) > 1 else v
+                            mm_payload[k] = element
+                    except Exception as e:
+                        # Best-effort; skip malformed entries
+                        logger.error(f"Error in merge multimodal outputs: {e}")
+                if mm_payload:
+                    payload.update(mm_payload)
+            pooler_output.append(payload)  # type: ignore[arg-type]
+            prev_logits_index = logits_index + 1
+
+        # Omni-new
+        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,
+        )
+        #  ------------------------------------------------------------------------------------------------
+
+        durations = ProfileExecuteDuration().pop_captured_sync()
+        if durations:
+            dr_str = [f"[{tag}]:{duration:.2f}ms" for tag, duration in durations.items()]
+            captured_name = "Decode" if self.attn_state == AscendAttentionState.DecodeOnly else "Prefill"
+            logger.info("Profile execute duration [%s]:%s", captured_name, " ".join(dr_str))
+        if self.dynamic_eplb:
+            self.eplb_updator.forward_end()
+        if not self.use_async_scheduling:
+            return output
+
+        return AsyncNPUModelRunnerOutput(
+            model_runner_output=output,
+            sampled_token_ids=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)
+
+    def _generate_process_reqs_hidden_states(
+        self,
+        attn_metadata,
+        with_prefill,
+        maybe_padded_num_tokens,
+        input_ids,
+        positions,
+        intermediate_tensors,
+        inputs_embeds,
+        model_kwargs_extra,
+    ):
+        assert self.model is not None
+        hidden_states = self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **model_kwargs_extra,
+        )
+
+        forward_context = get_forward_context()
+        if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL:
+            # TODO: maybe_padded_num_tokens will be removed, use num_input_tokens instead
+            if self.vllm_config.model_config.use_mla:
+                # FIXME: Try using `auto_dispatch_capture=True`
+                update_mla_attn_params(
+                    self.update_stream, forward_context, maybe_padded_num_tokens, self.speculative_config
+                )
+            else:
+                update_attn_params(self.update_stream, forward_context, maybe_padded_num_tokens)
+
+        if get_forward_context().sp_enabled:
+            hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
+            pad_size = get_forward_context().pad_size
+            if pad_size > 0:
+                hidden_states = hidden_states[:-pad_size, :]
+        return hidden_states

vllm_omni/worker/npu/npu_ar_worker.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_ascend.worker.worker_v1 import NPUWorker
+
+from vllm_omni.worker.npu.npu_ar_model_runner import NPUARModelRunner
+
+
+class NPUARWorker(NPUWorker):
+    """NPU AR worker for thinker/talker stages in Omni model."""
+
+    def init_device(self):
+        device = self._init_device()
+
+        self.model_runner: NPUARModelRunner = NPUARModelRunner(self.vllm_config, device)

vllm_omni/worker/npu/npu_diffusion_model_runner.py

@@ -0,0 +1,748 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from vllm.config import CUDAGraphMode
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.forward_context import BatchDescriptor
+from vllm.logger import logger
+from vllm.multimodal.inputs import MultiModalKwargs
+from vllm.sequence import IntermediateTensors
+from vllm.utils import (
+    cdiv,
+)
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
+from vllm.v1.core.sched.output import SchedulerOutput
+
+# yapf conflicts with isort for this block
+# yapf: disable
+from vllm.v1.kv_cache_interface import (
+    EncoderOnlyAttentionSpec,
+)
+
+# yapf: enable
+from vllm.v1.outputs import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    AsyncModelRunnerOutput,
+    ModelRunnerOutput,
+)
+from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
+from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput
+from vllm_ascend.ascend_forward_context import set_ascend_forward_context
+from vllm_ascend.attention.attention_v1 import AscendAttentionState
+from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
+from vllm_ascend.spec_decode.interface import SpecDcodeType
+from vllm_ascend.utils import (
+    ProfileExecuteDuration,
+    enable_sp,
+    lmhead_tp_enable,
+)
+from vllm_ascend.worker.model_runner_v1 import AsyncNPUModelRunnerOutput
+
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner
+
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.output import SchedulerOutput
+
+
+class NPUDiffusionModelRunner(OmniNPUModelRunner):
+    """Diffusion model runner for vLLM-omni on NPU (non-autoregressive)."""
+
+    def _prepare_inputs(
+        self,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+    ) -> tuple[
+        dict[str, Any],
+        torch.Tensor,
+        np.ndarray,
+        int,
+        torch.Tensor,
+        int,
+        torch.Tensor,
+        SpecDecodeMetadata,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        int,
+        dict[str, Any],
+    ]:
+        total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        assert total_num_scheduled_tokens > 0
+        num_reqs = self.input_batch.num_reqs
+        assert num_reqs > 0
+
+        # OPTIMIZATION: Start copying the block table first.
+        # This way, we can overlap the copy with the following CPU operations.
+        self.input_batch.block_table.commit_block_table(num_reqs)
+
+        # Get the number of scheduled tokens for each request.
+        req_ids = self.input_batch.req_ids
+        tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
+        num_scheduled_tokens = np.array(tokens, dtype=np.int32)
+        max_num_scheduled_tokens = num_scheduled_tokens.max()
+        num_valid_tokens = np.array(
+            [
+                num_tokens - len(scheduler_output.scheduled_spec_decode_tokens.get(i, []))
+                for num_tokens, i in zip(tokens, req_ids)
+            ],
+            dtype=np.int32,
+        )
+
+        if self.use_aclgraph and total_num_scheduled_tokens <= self.aclgraph_batch_sizes[-1]:
+            # Add padding to the batch size.
+            num_input_tokens = self.vllm_config.pad_for_cudagraph(total_num_scheduled_tokens)
+        elif self.use_aclgraph and enable_sp(self.vllm_config):
+            # When using aclgraph, if total_num_scheduled_tokens exceeds the maximum graph size,
+            # the model will fall back to running its FX graph in eager mode.
+            # In this case, when sequence parallelism is enabled, we need to pad tokens to align
+            # with tp_size because pad_size cannot be captured by the FX graph
+            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
+            num_input_tokens = math.ceil(total_num_scheduled_tokens / tp_size) * tp_size
+        else:
+            # Eager mode.
+            num_input_tokens = total_num_scheduled_tokens
+
+        # Get the attention state.
+        attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, num_valid_tokens)
+        self.attn_state = attn_state  # type: ignore
+
+        # Determine if it's a splitfuse batch
+        with_prefill = attn_state not in [AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding]
+
+        self.query_lens = torch.from_numpy(num_scheduled_tokens)
+        enable_dbo = self._check_dbo_is_valid(self.query_lens.tolist(), attn_state, total_num_scheduled_tokens)
+
+        # Get info across DP ranks.
+        # NOTE: maybe_padded_num_tokens is only used when using TorchAir with DP,
+        # Otherwise, it's just max_tokens_across_dp_cpu
+        (maybe_padded_num_tokens, num_tokens_across_dp, with_prefill, enable_dbo) = self._sync_metadata_across_dp(
+            num_input_tokens, with_prefill, enable_dbo
+        )
+
+        # TODO: Now that num_input_tokens is basically identical with maybe_padded_num_tokens
+        # We should consider removing maybe_padded_num_tokens later
+        num_input_tokens = maybe_padded_num_tokens
+
+        # Hot-Swap lora model
+        if self.lora_config:
+            self.set_active_loras(self.input_batch, num_scheduled_tokens)
+
+        # Get request indices.
+        # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
+        req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)
+
+        # cu_num_tokens: [2, 5, 3] -> [2, 7, 10]
+        # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
+        cu_num_tokens, arange = self._get_cumsum_and_arange(num_scheduled_tokens)
+
+        positions_np = self.positions_np[:total_num_scheduled_tokens]
+        np.add(self.input_batch.num_computed_tokens_cpu[req_indices], arange, out=positions_np)
+
+        # Calculate M-RoPE positions.
+        # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
+        if self.uses_mrope:
+            self._calc_mrope_positions(scheduler_output)
+
+            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
+            self.mrope_positions[:, :total_num_scheduled_tokens].copy_(
+                self.mrope_positions_cpu[:, :total_num_scheduled_tokens], non_blocking=True
+            )
+
+        # Get token indices.
+        # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
+        # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]
+        # where M is the max_model_len.
+        token_indices = positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]
+
+        # Prepare input_ids.
+        # NOTE(woosuk): We use torch.index_select instead of np.take here
+        # because torch.index_select is much faster than np.take for large
+        # tensors.
+        torch.index_select(
+            self.input_batch.token_ids_cpu_tensor.flatten(),
+            0,
+            torch.from_numpy(token_indices),
+            out=self.input_ids_cpu[:total_num_scheduled_tokens],
+        )
+
+        # Prepare some information for building Attention-Metadata
+        # Compute and commit slot mapping
+        self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np)
+        self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens)
+
+        self.query_start_loc_np[0] = 0
+        self.query_start_loc_np[1 : num_reqs + 1] = cu_num_tokens
+        self.query_start_loc[: num_reqs + 1].copy_(self.query_start_loc_cpu[: num_reqs + 1], non_blocking=True)
+
+        self.seq_lens_np[:num_reqs] = self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens
+        self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs], non_blocking=True)
+
+        # Fill unused with -1. Needed for reshape_and_cache
+        self.query_start_loc[num_reqs + 1 :].fill_(-1)
+        self.seq_lens[num_reqs:].fill_(0)
+
+        self.query_lens = torch.from_numpy(num_scheduled_tokens)
+
+        # Copy the tensors to the NPU.
+        self._prepare_input_ids(total_num_scheduled_tokens, cu_num_tokens)
+        self.positions_cpu[total_num_scheduled_tokens:num_input_tokens].zero_()
+        self.positions[:num_input_tokens].copy_(self.positions_cpu[:num_input_tokens], non_blocking=True)
+
+        # Make Attention metadata
+        positions_cpu = self.positions_cpu[:num_input_tokens]
+        positions = self.positions[:num_input_tokens]
+        seq_lens_cpu = self.seq_lens_cpu[:num_reqs]
+        attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, num_valid_tokens)
+        self.attn_mask = self._make_attention_mask(seq_lens=seq_lens_cpu, position=positions_cpu, attn_state=attn_state)
+        self.attn_state = attn_state  # type: ignore
+
+        self.with_prefill = with_prefill
+        self.num_tokens_across_dp = num_tokens_across_dp
+        self._update_graph_pad_size(with_prefill, maybe_padded_num_tokens)
+        attn_metadata: dict[str, Any] = {}
+
+        # _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)
+
+            # -------------------------------------- Omni-new -------------------------------------------------
+            # 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.
+            input_ids = self.input_ids[:num_input_tokens]
+            if mm_embeds:
+                inputs_embeds = self.model.get_input_embeddings(input_ids, mm_embeds)
+            else:
+                inputs_embeds = self.model.get_input_embeddings(input_ids)
+            # TODO(woosuk): Avoid the copy. Optimize.
+            self.inputs_embeds[:num_input_tokens].copy_(inputs_embeds)
+            inputs_embeds = self.inputs_embeds[:num_input_tokens]
+
+            model_kwargs = {
+                **self._init_model_kwargs(num_input_tokens),
+                **self._extract_mm_kwargs(scheduler_output),
+            }
+            # (NOTE) Omni-new: input_ids isn't set as 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 ACL graph.
+            input_ids = self.input_ids[:num_input_tokens]
+            inputs_embeds = None
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+        positions = self.positions[:num_input_tokens]
+        input_ids, positions = self._update_input_ids_and_positions(
+            input_ids, positions, num_input_tokens, with_prefill, maybe_padded_num_tokens
+        )
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            assert intermediate_tensors is not None
+            assert self.intermediate_tensors is not None
+            for k, v in intermediate_tensors.items():
+                self.intermediate_tensors[k][:num_input_tokens].copy_(v[:num_input_tokens], non_blocking=True)
+            intermediate_tensors = IntermediateTensors(
+                {k: v[:num_input_tokens] for k, v in self.intermediate_tensors.items()}
+            )
+
+        use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
+        if not use_spec_decode:
+            # NOTE(woosuk): Due to chunked prefills, the batch may contain
+            # partial requests. While we should not sample any token
+            # from these partial requests, we do so for simplicity.
+            # We will ignore the sampled tokens from the partial requests.
+            # TODO: Support prompt logprobs.
+            spec_decode_metadata = None
+            logits_indices = torch.from_numpy(cu_num_tokens - 1).to(self.device, non_blocking=True)
+        else:
+            # Get the number of draft tokens for each request.
+            # Iterate over the dictionary rather than all requests since not all
+            # requests have draft tokens.
+            num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
+            for req_id, draft_token_ids in scheduler_output.scheduled_spec_decode_tokens.items():
+                req_idx = self.input_batch.req_id_to_index[req_id]
+                num_draft_tokens[req_idx] = len(draft_token_ids)
+
+            spec_decode_metadata = self._calc_spec_decode_metadata(num_draft_tokens, cu_num_tokens)
+            logits_indices = spec_decode_metadata.logits_indices
+            self.num_draft_tokens.np[:num_reqs] = num_draft_tokens
+            self.num_draft_tokens.np[num_reqs:].fill(0)
+            self.num_draft_tokens.copy_to_gpu()
+
+        # Used in the below loop.
+        # query_start_loc_cpu = self.query_start_loc.cpu[:num_reqs + 1]
+        num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs]
+        spec_decode_common_attn_metadata = None
+        if use_spec_decode and self.need_accepted_tokens:
+            self.num_accepted_tokens.np[:num_reqs] = self.input_batch.num_accepted_tokens_cpu[:num_reqs]
+            self.num_accepted_tokens.np[num_reqs:].fill(1)
+            self.num_accepted_tokens.copy_to_gpu()
+
+        # Prepare the attention metadata for each KV cache group and make layers
+        # in the same group share the same metadata.
+        for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups):
+            if isinstance(kv_cache_group_spec.kv_cache_spec, EncoderOnlyAttentionSpec):
+                # Encoder-only layers do not have KV cache, so we need to
+                # create a dummy block table and slot mapping for them.
+                blk_table_tensor = torch.zeros(
+                    (num_reqs, 1),
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+                slot_mapping = torch.zeros(
+                    (total_num_scheduled_tokens,),
+                    dtype=torch.int64,
+                    device=self.device,
+                )
+            else:
+                blk_table = self.input_batch.block_table[kv_cache_group_id]
+                blk_table_tensor = blk_table.get_device_tensor()
+                slot_mapping = blk_table.slot_mapping_cpu[:total_num_scheduled_tokens]
+                self.slot_mapping[:total_num_scheduled_tokens].copy_(
+                    slot_mapping[:total_num_scheduled_tokens],
+                    non_blocking=True,
+                )
+                self.slot_mapping[total_num_scheduled_tokens:].fill_(0)
+
+            # Make AscendCommonAttentionMetadata
+            common_attn_metadata = AscendCommonAttentionMetadata(
+                query_start_loc=self.query_start_loc[: num_reqs + 1],
+                query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1],
+                seq_lens_cpu=self.seq_lens_cpu,
+                seq_lens=self.seq_lens_cpu[:num_reqs],
+                num_reqs=num_reqs,
+                num_actual_tokens=total_num_scheduled_tokens,
+                num_input_tokens=num_input_tokens,
+                actual_seq_lengths_q=self.actual_seq_lengths_q,
+                # TODO: change this to the right block table for linear attn
+                block_table_tensor=blk_table_tensor[:num_reqs],
+                slot_mapping=self.slot_mapping,
+                num_computed_tokens_cpu=num_computed_tokens_cpu,
+                positions=self.positions,
+                attn_mask=self.attn_mask,
+                spec_attn_mask=self.spec_attn_mask,
+                attn_state=self.attn_state,
+                enable_dbo_across_dp=enable_dbo,
+                is_only_prefill=bool(np.all(num_valid_tokens != 1)),
+                max_query_len=max_num_scheduled_tokens,
+                graph_pad_size=self.graph_pad_size,
+                decode_token_per_req=self.decode_token_per_req,
+                cos=self.cos,
+                sin=self.sin,
+            )
+
+            if self.speculative_config and spec_decode_common_attn_metadata is None:
+                spec_decode_common_attn_metadata = common_attn_metadata
+
+            for attn_group in self.attn_groups[kv_cache_group_id]:
+                common_prefix_len = 0
+                extra_attn_metadata_args = {}
+                builder = attn_group.get_metadata_builder()
+                if isinstance(builder, GDNAttentionMetadataBuilder) or self.model_config.runner_type == "pooling":
+                    if use_spec_decode:
+                        extra_attn_metadata_args = dict(
+                            num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs],
+                            num_draft_tokens=self.num_draft_tokens.gpu[:num_reqs],
+                        )
+                    attn_metadata_i = builder.build(
+                        common_prefix_len=common_prefix_len,
+                        common_attn_metadata=common_attn_metadata,
+                        **extra_attn_metadata_args,
+                    )
+                else:
+                    attn_metadata_i = builder.build(
+                        common_prefix_len=common_prefix_len,
+                        common_attn_metadata=common_attn_metadata,
+                        model=self.get_model(),
+                        **extra_attn_metadata_args,
+                    )
+
+                for layer_name in attn_group.layer_names:
+                    attn_metadata[layer_name] = attn_metadata_i
+
+        if lmhead_tp_enable():
+            max_num_reqs_across_dp = maybe_padded_num_tokens if not with_prefill else self.max_num_reqs
+            logits_indices = nn.functional.pad(logits_indices, (0, max_num_reqs_across_dp - logits_indices.shape[0]))
+
+        return (
+            attn_metadata,
+            positions,
+            num_input_tokens,
+            num_input_tokens,
+            num_tokens_across_dp,
+            maybe_padded_num_tokens,
+            logits_indices,
+            spec_decode_metadata,
+            input_ids,
+            inputs_embeds,
+            intermediate_tensors,
+            max_num_scheduled_tokens,
+            model_kwargs,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+    ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
+        with ProfileExecuteDuration().capture_async("prepare input"):
+            self._update_states(scheduler_output)
+            if not scheduler_output.total_num_scheduled_tokens:
+                return EMPTY_MODEL_RUNNER_OUTPUT
+
+            if self.dynamic_eplb:
+                self.eplb_updator.forward_before()
+
+            (
+                attn_metadata,
+                positions,
+                num_scheduled_tokens_np,
+                num_input_tokens,
+                num_tokens_across_dp,
+                maybe_padded_num_tokens,
+                logits_indices,
+                spec_decode_metadata,
+                input_ids,
+                inputs_embeds,
+                intermediate_tensors,
+                max_query_len,
+                model_kwargs,
+            ) = self._prepare_inputs(scheduler_output, intermediate_tensors)
+
+            if self.dynamic_eplb:
+                self.eplb_updator.take_update_info_from_eplb_process()
+
+        moe_comm_type = self._select_moe_comm_method(num_input_tokens, self.with_prefill)
+
+        # -------------------------------------- Omni-new -------------------------------------------------
+        # Omni-new: don't use cudagraph_dispatcher
+        # and remove ubatch_slices
+        aclgraph_runtime_mode = CUDAGraphMode.NONE
+        # -------------------------------------------------------------------------------------------------
+
+        # Run forward pass
+        with ProfileExecuteDuration().capture_async("forward"):
+            with set_ascend_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                with_prefill=self.with_prefill,
+                reserved_mc2_mask=self.reserved_mc2_mask,
+                moe_comm_type=moe_comm_type,
+                aclgraph_runtime_mode=aclgraph_runtime_mode,
+                batch_descriptor=None,
+                num_actual_tokens=scheduler_output.total_num_scheduled_tokens,
+                prefetch_stream=self.prefetch_stream,
+                model_instance=self.model,
+                weight_prefetch_method=self.weight_prefetch_method,
+            ):
+                self.maybe_setup_kv_connector(scheduler_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,
+                )
+
+            self.maybe_wait_for_kv_save()
+            finished_sending, finished_recving = self.get_finished_kv_transfer(scheduler_output)
+
+            aux_hidden_states = None
+            if self.drafter and self.drafter.name == SpecDcodeType.EAGLE3:
+                hidden_states, aux_hidden_states = outputs
+
+        kv_connector_output = KVConnectorOutput(finished_sending=finished_sending, finished_recving=finished_recving)
+        finished_sending = None
+        finished_recving = None
+
+        # -------------------------------------- Omni-new -------------------------------------------------
+        # Omni-new: extract_multimodal_outputs
+        _, multimodal_outputs = self.extract_multimodal_outputs(outputs)
+        # Ensure one tensor per request, map to CPU for output struct
+        pooler_output: list[torch.Tensor | None] = []
+        if isinstance(multimodal_outputs, torch.Tensor):
+            # If model returned a single stacked tensor, split by requests
+            assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
+            for i in range(self.input_batch.num_reqs):
+                pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous())
+        elif isinstance(multimodal_outputs, list):
+            for out in multimodal_outputs:
+                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+        elif isinstance(multimodal_outputs, dict):
+            for out in multimodal_outputs.values():
+                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+        else:
+            raise RuntimeError("Unsupported diffusion output type")
+
+        output = OmniModelRunnerOutput(
+            req_ids=self.input_batch.req_ids,
+            req_id_to_index=self.input_batch.req_id_to_index,
+            sampled_token_ids=[],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=pooler_output,
+            kv_connector_output=kv_connector_output,
+            num_nans_in_logits={},
+        )
+        # -------------------------------------------------------------------------------------------------
+
+        durations = ProfileExecuteDuration().pop_captured_sync()
+        if durations:
+            dr_str = [f"[{tag}]:{duration:.2f}ms" for tag, duration in durations.items()]
+            captured_name = "Decode" if self.attn_state == AscendAttentionState.DecodeOnly else "Prefill"
+            logger.info("Profile execute duration [%s]:%s", captured_name, " ".join(dr_str))
+
+        if self.dynamic_eplb:
+            self.eplb_updator.forward_end()
+
+        if not self.use_async_scheduling:
+            return output
+
+        return AsyncNPUModelRunnerOutput(
+            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: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None,
+        multimodal_kwargs: dict,
+        logits_indices: torch.Tensor,
+    ) -> 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)
+        """
+        kwargs = dict(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
+            sampling_metadata=self.input_batch.sampling_metadata,
+            logits_index=logits_indices,
+            sampler=self.sampler,
+        )
+
+        if hasattr(self.model, "forward"):
+            return self.model.forward(**kwargs)
+        # 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,
+        with_prefill: bool = False,
+        is_torchair_compile: bool = False,
+        aclgraph_runtime_mode: CUDAGraphMode | None = None,
+        force_attention: bool = False,
+        uniform_decode: bool = False,
+    ) -> torch.Tensor:
+        # only support eager mode and piecewise graph now
+        assert aclgraph_runtime_mode is None or aclgraph_runtime_mode in {
+            CUDAGraphMode.NONE,
+            CUDAGraphMode.PIECEWISE,
+            CUDAGraphMode.FULL,
+        }
+
+        # In multi-DP scenarios, there may be situations where all DP groups are executing dummy runs.
+        # If sequence parallelism is enabled, it is essential to ensure that num_tokens is divisible by tp_size.
+        if self.use_aclgraph and enable_sp(self.vllm_config):
+            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
+            num_tokens = math.ceil(num_tokens / tp_size) * tp_size
+
+        # Force dummy run on prefill stage when this node is deemed as kv producer.
+        if self.is_kv_producer and not self.is_kv_consumer:
+            with_prefill = True
+
+        # Padding for DP
+        (num_tokens, num_tokens_across_dp, with_prefill, _) = self._sync_metadata_across_dp(
+            num_tokens, with_prefill, False
+        )
+
+        moe_comm_type = self._select_moe_comm_method(num_tokens, with_prefill)
+
+        # 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.max_num_reqs
+        if uniform_decode:
+            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:
+            if with_prefill:
+                num_reqs = num_tokens
+            else:
+                num_reqs = (num_tokens + self.decode_token_per_req - 1) // self.decode_token_per_req
+            num_reqs = min(num_reqs, 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)
+
+        if not self.in_profile_run and self.dynamic_eplb:
+            self.eplb_updator.forward_before()
+
+        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]
+            else:
+                input_ids = self.input_ids[:num_tokens]
+                inputs_embeds = None
+
+            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=num_tokens, dtype=self.dtype, device=self.device
+                    )
+                intermediate_tensors = IntermediateTensors(
+                    {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()}
+                )
+
+            # filter out the valid batch descriptor
+            _ag_mode, batch_descriptor = self.aclgraph_dispatcher.dispatch(
+                BatchDescriptor(num_tokens=num_tokens, uniform_decode=uniform_decode)
+            )
+            if aclgraph_runtime_mode is not None:
+                # we allow forcing NONE when the dispatcher disagrees to support
+                # warm ups for aclgraph capture
+                assert aclgraph_runtime_mode == CUDAGraphMode.NONE or aclgraph_runtime_mode == _ag_mode, (
+                    f"Aclgraph runtime mode mismatch at dummy_run. "
+                    f"Expected {_ag_mode}, but got {aclgraph_runtime_mode}."
+                )
+            else:
+                aclgraph_runtime_mode = _ag_mode
+
+            # TODO(Mengqing): Set create_mixed_batch to False since it's only used in FI warmup
+            # and not supported in ASCEND now. We could remove it in the future.
+            attn_metadata = self._build_dummy_attn_metadata(
+                False,
+                num_reqs=num_reqs,
+                num_tokens=num_tokens,
+                max_query_len=max_query_len,
+                aclgraph_runtime_mode=aclgraph_runtime_mode,
+                force_attention=force_attention,
+            )
+
+            need_dummy_logits = not self.in_profile_run and lmhead_tp_enable()
+
+            if need_dummy_logits:
+                max_num_reqs_across_dp = num_tokens if not with_prefill else max_num_reqs
+                dummy_indices = torch.zeros(max_num_reqs_across_dp, dtype=torch.int32)
+
+                def dummy_compute_logits(hidden_states):
+                    return self.model.compute_logits(hidden_states[dummy_indices])
+
+            with set_ascend_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                with_prefill=with_prefill,
+                in_profile_run=self.in_profile_run,
+                reserved_mc2_mask=self.reserved_mc2_mask,
+                moe_comm_type=moe_comm_type,
+                num_actual_tokens=0,
+                aclgraph_runtime_mode=aclgraph_runtime_mode,
+                batch_descriptor=batch_descriptor,
+                prefetch_stream=self.prefetch_stream,
+                model_instance=self.model,
+                weight_prefetch_method=self.weight_prefetch_method,
+            ):
+                hidden_states = self._generate_dummy_run_hidden_states(
+                    with_prefill,
+                    is_torchair_compile,
+                    input_ids,
+                    positions,
+                    attn_metadata,
+                    num_tokens,
+                    intermediate_tensors,
+                    inputs_embeds,
+                )
+                if need_dummy_logits:
+                    dummy_compute_logits(hidden_states)
+
+            if self.drafter:
+                self.drafter.dummy_run(
+                    num_tokens=num_tokens,
+                    with_prefill=with_prefill,
+                    skip_attn=True,
+                    num_reqs=num_reqs,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    aclgraph_runtime_mode=aclgraph_runtime_mode,
+                    batch_descriptor=batch_descriptor,
+                )
+                if need_dummy_logits:
+                    self.drafter.model.compute_logits(hidden_states[dummy_indices])
+            if self.in_profile_run and self.dynamic_eplb:
+                self.model.clear_all_moe_loads()
+            if not self.in_profile_run and self.dynamic_eplb:
+                self.eplb_updator.take_update_info_from_eplb_process()
+                self.eplb_updator.forward_end()
+
+            # -------------------------------------- Omni-new -------------------------------------------------
+            hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
+            # -------------------------------------------------------------------------------------------------
+            return hidden_states

vllm_omni/worker/npu/npu_diffusion_worker.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_ascend.worker.worker_v1 import NPUWorker
+
+from vllm_omni.worker.npu.npu_diffusion_model_runner import NPUDiffusionModelRunner
+
+
+class NPUDiffusionWorker(NPUWorker):
+    """NPU diffusion worker for code2wav stage in Omni model."""
+
+    def init_device(self):
+        device = self._init_device()
+
+        self.model_runner: NPUDiffusionModelRunner = NPUDiffusionModelRunner(self.vllm_config, device)

vllm_omni/worker/npu/npu_model_runner.py

@@ -0,0 +1,543 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+from typing import TYPE_CHECKING, Any, 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
+from vllm.logger import init_logger
+from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
+from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_mrope
+from vllm.model_executor.models.interfaces_base import VllmModelForPooling
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import BatchedTensorInputs, MultiModalKwargsItem
+from vllm.multimodal.utils import group_mm_kwargs_by_modality
+from vllm.sampling_params import SamplingType
+from vllm.utils import cdiv
+from vllm.v1.worker.gpu_model_runner import IntermediateTensors
+from vllm_ascend.ascend_forward_context import set_ascend_forward_context
+from vllm_ascend.utils import enable_sp, lmhead_tp_enable
+from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
+from vllm_ascend.worker.npu_input_batch import CachedRequestState
+
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.output import SchedulerOutput
+
+logger = init_logger(__name__)
+
+
+class OmniNPUModelRunner(NPUModelRunner):
+    """
+    Base class for NPU model runners with multimodality support.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.is_multimodal_raw_input_only_model = self.model_config.is_multimodal_raw_input_only_model
+        self.mm_registry = MULTIMODAL_REGISTRY
+        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(self.model_config)
+
+    def _init_mrope_positions(self, req_state: CachedRequestState):
+        image_grid_thw = []
+        video_grid_thw = []
+        second_per_grid_ts = []
+        audio_feature_lengths = []
+        use_audio_in_video = False
+        for mm_feature in req_state.mm_features:
+            mm_item = mm_feature.data
+            if mm_item is None:
+                continue
+            mm_input = mm_item.get_data()
+            if (t := mm_input.get("image_grid_thw")) is not None:
+                image_grid_thw.append(t.tolist())
+            if (t := mm_input.get("video_grid_thw")) is not None:
+                video_grid_thw.append(t.tolist())
+            if (t := mm_input.get("second_per_grid_ts")) is not None:
+                second_per_grid_ts.append(t)
+            if (t := mm_input.get("audio_feature_lengths")) is not None:
+                audio_feature_lengths.append(t)
+            # Check for use_audio_in_video
+            use_audio_in_video_value = mm_input.get("use_audio_in_video")
+            if use_audio_in_video_value is not None:
+                use_audio_in_video = bool(use_audio_in_video_value.item())
+
+        assert supports_mrope(self.get_model()), "M-RoPE support is not implemented."
+
+        req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor(
+            req_state.prompt_token_ids,
+            hf_config=self.model_config.hf_config,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            second_per_grid_ts=second_per_grid_ts,
+            audio_feature_lengths=audio_feature_lengths,
+            use_audio_in_video=use_audio_in_video,
+        )
+
+    def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
+        # 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)
+        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)
+
+        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.get_model())
+                to_update = model.pooler.get_pooling_updates(task)
+                to_update.apply(pooling_params)
+
+            backward_kwargs = {}
+            backward_kwargs["mm_features"] = new_req_data.mm_features
+
+            self.requests[req_id] = CachedRequestState(
+                req_id=req_id,
+                prompt_token_ids=new_req_data.prompt_token_ids,
+                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,
+                **backward_kwargs,
+            )
+
+            #  -------------------------------------- Omni-new -------------------------------------------------
+            # 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(self.requests[req_id])
+
+            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:
+                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.
+                req_ids_to_add.append(req_id)
+                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 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 extract_multimodal_outputs(
+        self, hidden_states: Union[torch.Tensor, list[torch.Tensor]]
+    ) -> tuple[torch.Tensor, Union[torch.Tensor, list[torch.Tensor], dict]]:
+        """Extract multimodal outputs from hidden states."""
+        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
+
+    def _init_model_kwargs(self, num_tokens: int):
+        model_kwargs = dict[str, Any]()
+
+        if not self.is_pooling_model:
+            return model_kwargs
+
+        num_reqs = self.input_batch.num_reqs
+        pooling_params = self.input_batch.get_pooling_params()
+
+        token_type_id_requests = dict[int, Any]()
+        for i, param in enumerate(pooling_params):
+            if (
+                param.extra_kwargs is not None
+                and (token_types := param.extra_kwargs.get("compressed_token_type_ids")) is not None
+            ):
+                token_type_id_requests[i] = token_types
+
+        if len(token_type_id_requests) == 0:
+            return model_kwargs
+
+        seq_lens = self.seq_lens_cpu[:num_reqs]
+        token_type_ids = []
+
+        for i in range(num_reqs):
+            pos = token_type_id_requests.get(i, seq_lens[i])
+            ids = (torch.arange(seq_lens[i]) >= pos).int()
+            token_type_ids.append(ids)
+
+        model_kwargs["token_type_ids"] = torch.concat(token_type_ids).to(device=self.device)
+        return model_kwargs
+
+    def _extract_mm_kwargs(
+        self,
+        scheduler_output: "SchedulerOutput",
+    ) -> BatchedTensorInputs:
+        if not scheduler_output or not self.is_multimodal_raw_input_only_model:
+            return {}
+
+        mm_kwargs = list[MultiModalKwargsItem]()
+        for req in scheduler_output.scheduled_new_reqs:
+            for feature in req.mm_features:
+                if feature.data is not None:
+                    mm_kwargs.append(feature.data)
+
+        # Input all modalities at once
+        model = cast(SupportsMultiModal, self.model)
+        mm_kwargs_combined: BatchedTensorInputs = {}
+        for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(
+            mm_kwargs,
+            device=self.device,
+            pin_memory=self.pin_memory,
+            merge_by_field_config=model.merge_by_field_config,
+        ):
+            mm_kwargs_combined.update(mm_kwargs_group)
+
+        return mm_kwargs_combined
+
+    @torch.inference_mode()
+    def _dummy_run(
+        self,
+        num_tokens: int,
+        with_prefill: bool = False,
+        is_torchair_compile: bool = False,
+        aclgraph_runtime_mode: Optional[CUDAGraphMode] = None,
+        force_attention: bool = False,
+        uniform_decode: bool = False,
+    ) -> torch.Tensor:
+        # only support eager mode and piecewise graph now
+        assert aclgraph_runtime_mode is None or aclgraph_runtime_mode in {
+            CUDAGraphMode.NONE,
+            CUDAGraphMode.PIECEWISE,
+            CUDAGraphMode.FULL,
+        }
+
+        # In multi-DP scenarios, there may be situations where all DP groups are executing dummy runs.
+        # If sequence parallelism is enabled, it is essential to ensure that num_tokens is divisible by tp_size.
+        if self.use_aclgraph and enable_sp(self.vllm_config):
+            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
+            num_tokens = math.ceil(num_tokens / tp_size) * tp_size
+
+        # Force dummy run on prefill stage when this node is deemed as kv producer.
+        if self.is_kv_producer and not self.is_kv_consumer:
+            with_prefill = True
+
+        # Padding for DP
+        (num_tokens, num_tokens_across_dp, with_prefill, _) = self._sync_metadata_across_dp(
+            num_tokens, with_prefill, False
+        )
+
+        moe_comm_type = self._select_moe_comm_method(num_tokens, with_prefill)
+
+        # 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.max_num_reqs
+        if uniform_decode:
+            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:
+            if with_prefill:
+                num_reqs = num_tokens
+            else:
+                num_reqs = (num_tokens + self.decode_token_per_req - 1) // self.decode_token_per_req
+            num_reqs = min(num_reqs, 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)
+
+        if not self.in_profile_run and self.dynamic_eplb:
+            self.eplb_updator.forward_before()
+
+        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]
+            else:
+                input_ids = self.input_ids[:num_tokens]
+                inputs_embeds = None
+
+            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=num_tokens, dtype=self.dtype, device=self.device
+                    )
+                intermediate_tensors = IntermediateTensors(
+                    {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()}
+                )
+
+            # filter out the valid batch descriptor
+            _ag_mode, batch_descriptor = self.aclgraph_dispatcher.dispatch(
+                BatchDescriptor(num_tokens=num_tokens, uniform_decode=uniform_decode)
+            )
+            if aclgraph_runtime_mode is not None:
+                # we allow forcing NONE when the dispatcher disagrees to support
+                # warm ups for aclgraph capture
+                assert aclgraph_runtime_mode == CUDAGraphMode.NONE or aclgraph_runtime_mode == _ag_mode, (
+                    f"Aclgraph runtime mode mismatch at dummy_run. "
+                    f"Expected {_ag_mode}, but got {aclgraph_runtime_mode}."
+                )
+            else:
+                aclgraph_runtime_mode = _ag_mode
+
+            # TODO(Mengqing): Set create_mixed_batch to False since it's only used in FI warmup
+            # and not supported in ASCEND now. We could remove it in the future.
+            attn_metadata = self._build_dummy_attn_metadata(
+                False,
+                num_reqs=num_reqs,
+                num_tokens=num_tokens,
+                max_query_len=max_query_len,
+                aclgraph_runtime_mode=aclgraph_runtime_mode,
+                force_attention=force_attention,
+            )
+
+            need_dummy_logits = not self.in_profile_run and lmhead_tp_enable()
+
+            if need_dummy_logits:
+                max_num_reqs_across_dp = num_tokens if not with_prefill else max_num_reqs
+                dummy_indices = torch.zeros(max_num_reqs_across_dp, dtype=torch.int32)
+
+                def dummy_compute_logits(hidden_states):
+                    return self.model.compute_logits(hidden_states[dummy_indices])
+
+            with set_ascend_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                with_prefill=with_prefill,
+                in_profile_run=self.in_profile_run,
+                reserved_mc2_mask=self.reserved_mc2_mask,
+                moe_comm_type=moe_comm_type,
+                num_actual_tokens=0,
+                aclgraph_runtime_mode=aclgraph_runtime_mode,
+                batch_descriptor=batch_descriptor,
+                prefetch_stream=self.prefetch_stream,
+                model_instance=self.model,
+                weight_prefetch_method=self.weight_prefetch_method,
+            ):
+                hidden_states = self._generate_dummy_run_hidden_states(
+                    with_prefill,
+                    is_torchair_compile,
+                    input_ids,
+                    positions,
+                    attn_metadata,
+                    num_tokens,
+                    intermediate_tensors,
+                    inputs_embeds,
+                )
+                if need_dummy_logits:
+                    dummy_compute_logits(hidden_states)
+
+            if self.drafter:
+                self.drafter.dummy_run(
+                    num_tokens=num_tokens,
+                    with_prefill=with_prefill,
+                    skip_attn=True,
+                    num_reqs=num_reqs,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    aclgraph_runtime_mode=aclgraph_runtime_mode,
+                    batch_descriptor=batch_descriptor,
+                )
+                if need_dummy_logits:
+                    self.drafter.model.compute_logits(hidden_states[dummy_indices])
+            if self.in_profile_run and self.dynamic_eplb:
+                self.model.clear_all_moe_loads()
+            if not self.in_profile_run and self.dynamic_eplb:
+                self.eplb_updator.take_update_info_from_eplb_process()
+                self.eplb_updator.forward_end()
+
+            hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
+
+            return hidden_states