update

.gitignore

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

docs/.nav.yml

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

docs/api/README.md

@@ -4,8 +4,8 @@
 
 Main entry points for vLLM-Omni inference and serving.
 
-- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
-- [vllm_omni.entrypoints.async_omni.AsyncOmniStageLLM][]
+- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
+- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniStageLLM][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
 - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
@@ -34,8 +34,6 @@ Engine classes for offline and online inference.
 - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
-- [vllm_omni.engine.OmniEngineCoreOutput][]
-- [vllm_omni.engine.OmniEngineCoreOutputs][]
 - [vllm_omni.engine.OmniEngineCoreRequest][]
 - [vllm_omni.engine.PromptEmbedsPayload][]
 - [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
@@ -50,7 +48,6 @@ 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][]
 
@@ -58,7 +55,7 @@ Core scheduling and caching components.
 
 Model execution components.
 
-- [vllm_omni.model_executor.models.output_templates.OmniOutput][]
+- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.OmniOutput][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniAudioFeatureInputs][]
@@ -75,22 +72,6 @@ Model execution components.
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2EmbeddingModel][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2ForCausalLM][]
 - [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2Model][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_moe.Qwen3MoeForCausalLM][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni.Qwen3OmniMoeForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav.Qwen3OmniMoeCode2Wav][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniCodePredictorBaseModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniMoeTalkerCodePredictor][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMForCausalLM][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeConditionalGenerationMixin][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerMultiModalProcessor][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerProcessingInfo][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3Omni_VisionTransformer][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3_VisionPatchEmbed][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3_VisionPatchMerger][]
 
 ## Configuration
 
@@ -102,15 +83,10 @@ 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 AsyncOmni will use its tokenizer as default
+    is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmniLLM will use its tokenizer as default
     final_output: true # If the stage has output as part of final outputs. If it is false, which means that the stage only works as a intermediate role.
     final_output_type: text # What is the final output type. It can be text and audio now.
     default_sampling_params: # sampling parameters for the stage. Their meaning aligns with vLLM.
@@ -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, `AsyncOmni` will use its tokenizer as the default tokenizer for the entire pipeline.
+Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmniLLM` will use its tokenizer as the default tokenizer for the entire pipeline.
 
 Default: `true`
 

docs/contributing/DOCS_GUIDE.md

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

docs/design/api_design_doc.md

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

docs/design/api_design_template.md

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

docs/design/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/AsyncOmni (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
+API Server --> OmniLLM/AsyncOmniLLM (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
  --> Scheduler (New one for DiT) --> Executor (New one for diffusers) --> Worker (New one for DiT)
  --> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> 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 AsyncOmni]
+    D --> F[Initialize AsyncOmniLLM]
     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 AsyncOmni(AsyncLLM):
+class AsyncOmniLLM(AsyncLLM):
     """Extended AsyncLLM supporting multiple engines and stage-based processing"""
 
     def __init__(self, stage_configs: List[StageConfig]):
@@ -449,15 +449,15 @@ AsyncLLM.add_request()
 ↓
 OutputProcessor.add_request()
 ↓
-RequestState.from_new_request() → create request state
+RequestState.from_new_request() → 创建请求状态
 
-AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() →  create a background loop that continuously pulls output from EngineCore
+AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() →  创建一个Background loop 持续从EngineCore获取输出
 ↓
-OutputProcessor.process_outputs() → update state and process outputs
+OutputProcessor.process_outputs() → 更新状态并处理输出
 ↓
-RequestState.make_request_output() → convert to the final output, formatted as RequestOutput or PoolingRequestOutput
+RequestState.make_request_output() → 转换为最终输出,格式为RequestOutput或者PoolingRequestOutput
 ↓
-RequestOutputCollector.put() → push to the queue (AsyncLLM)
+RequestOutputCollector.put() → 推送到队列（AsyncLLM）
 ```
 
 Need to add implementation for an existing method

docs/getting_started/installation/README.md

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

docs/getting_started/installation/npu.md

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

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

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

docs/mkdocs/hooks/generate_examples.py

@@ -1,335 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import itertools
-import logging
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Literal
-
-import regex as re
-import yaml
-
-logger = logging.getLogger("mkdocs")
-
-ROOT_DIR = Path(__file__).parent.parent.parent.parent
-ROOT_DIR_RELATIVE = "../../../../.."
-EXAMPLE_DIR = ROOT_DIR / "examples"
-EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
-NAV_FILE = ROOT_DIR / "docs/.nav.yml"
-
-
-def fix_case(text: str) -> str:
-    subs = {
-        "api": "API",
-        "cli": "CLI",
-        "cpu": "CPU",
-        "llm": "LLM",
-        "mae": "MAE",
-        "tpu": "TPU",
-        "gguf": "GGUF",
-        "lora": "LoRA",
-        "rlhf": "RLHF",
-        "vllm": "vLLM",
-        "openai": "OpenAI",
-        "lmcache": "LMCache",
-        "multilora": "MultiLoRA",
-        "mlpspeculator": "MLPSpeculator",
-        r"fp\d+": lambda x: x.group(0).upper(),  # e.g. fp16, fp32
-        r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
-    }
-    for pattern, repl in subs.items():
-        text = re.sub(rf"\b{pattern}\b", repl, text, flags=re.IGNORECASE)
-    return text
-
-
-@dataclass
-class Example:
-    """
-    Example class for generating documentation content from a given path.
-
-    Attributes:
-        path (Path): The path to the main directory or file.
-        category (str): The category of the document.
-        main_file (Path): The main file in the directory.
-        other_files (list[Path]): list of other files in the directory.
-        title (str): The title of the document.
-
-    Methods:
-        __post_init__(): Initializes the main_file, other_files, and title attributes.
-        determine_main_file() -> Path: Determines the main file in the given path.
-        determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
-        determine_title() -> str: Determines the title of the document.
-        generate() -> str: Generates the documentation content.
-    """  # noqa: E501
-
-    path: Path
-    category: str = None
-    main_file: Path = field(init=False)
-    other_files: list[Path] = field(init=False)
-    title: str = field(init=False)
-
-    def __post_init__(self):
-        self.main_file = self.determine_main_file()
-        self.other_files = self.determine_other_files()
-        self.title = self.determine_title()
-
-    @property
-    def is_code(self) -> bool:
-        return self.main_file.suffix != ".md"
-
-    def determine_main_file(self) -> Path:
-        """
-        Determines the main file in the given path.
-        If the path is a file, it returns the path itself. Otherwise, it searches
-        for Markdown files (*.md) in the directory and returns the first one found.
-        Returns:
-            Path: The main file path, either the original path if it's a file or the first
-            Markdown file found in the directory.
-        Raises:
-            IndexError: If no Markdown files are found in the directory.
-        """  # noqa: E501
-        return self.path if self.path.is_file() else list(self.path.glob("*.md")).pop()
-
-    def determine_other_files(self) -> list[Path]:
-        """
-        Determine other files in the directory excluding the main file.
-
-        This method checks if the given path is a file. If it is, it returns an empty list.
-        Otherwise, it recursively searches through the directory and returns a list of all
-        files that are not the main file.
-
-        Returns:
-            list[Path]: A list of Path objects representing the other files in the directory.
-        """  # noqa: E501
-        if self.path.is_file():
-            return []
-        # Binary file extensions to exclude
-        binary_extensions = {
-            ".wav",
-            ".mp3",
-            ".mp4",
-            ".avi",
-            ".mov",
-            ".mkv",  # Audio/Video
-            ".png",
-            ".jpg",
-            ".jpeg",
-            ".gif",
-            ".bmp",
-            ".ico",
-            ".svg",  # Images
-            ".pdf",
-            ".zip",
-            ".tar",
-            ".gz",
-            ".bz2",
-            ".xz",  # Archives/Documents
-            ".exe",
-            ".so",
-            ".dll",
-            ".dylib",  # Binaries
-            ".bin",
-            ".dat",
-            ".db",
-            ".sqlite",  # Data files
-        }
-
-        def is_other_file(file: Path) -> bool:
-            return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions
-
-        return [file for file in self.path.rglob("*") if is_other_file(file)]
-
-    def determine_title(self) -> str:
-        if not self.is_code:
-            # Specify encoding for building on Windows
-            with open(self.main_file, encoding="utf-8") as f:
-                first_line = f.readline().strip()
-            match = re.match(r"^#\s+(?P<title>.+)$", first_line)
-            if match:
-                return match.group("title")
-        return fix_case(self.path.stem.replace("_", " ").title())
-
-    def fix_relative_links(self, content: str) -> str:
-        """
-        Fix relative links in markdown content by converting them to gh-file
-        format.
-
-        Args:
-            content (str): The markdown content to process
-
-        Returns:
-            str: Content with relative links converted to gh-file format
-        """
-        # Regex to match markdown links [text](relative_path)
-        # This matches links that don't start with http, https, ftp, or #
-        link_pattern = r"\[([^\]]*)\]\((?!(?:https?|ftp)://|#)([^)]+)\)"
-
-        def replace_link(match):
-            link_text = match.group(1)
-            relative_path = match.group(2)
-
-            # Make relative to repo root
-            gh_file = (self.main_file.parent / relative_path).resolve()
-            gh_file = gh_file.relative_to(ROOT_DIR)
-
-            # Make GitHub URL
-            url = "https://github.com/vllm-project/vllm/"
-            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,7 +17,6 @@ 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` |
 
@@ -34,4 +33,3 @@ th {
 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
-| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -1,40 +1,55 @@
-# Offline Example of vLLM-Omni for Qwen2.5-omni
+# Offline Inference Example of vLLM-Omni for Qwen2.5-Omni
 
-Source <https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen2_5_omni>.
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni>.
 
 
 ## 🛠️ Installation
 
-Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/README.md)
+Please refer to [installation](../../../getting_started/installation/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). To get the prompt, you can:
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit) and extract the prompts with `extract_prompts.py`
+??? abstract "extract_prompts.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
+    ``````
 ```bash
 tar -xf <Your Download Path>/seedtts_testset.tar
 cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
 python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
   --input examples/offline_inference/qwen2_5_omni/meta.lst \
   --output examples/offline_inference/qwen2_5_omni/top100.txt \
-  --topk 100
-```
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
+  --topk 10
 ```
 Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
-### Single Prompt
-Get into the example folder
-```bash
-cd examples/offline_inference/qwen2_5_omni
-```
-Then run the command below.
-```bash
-bash run_single_prompt.sh
-```
+??? abstract "run_multiple_prompts.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
+    ``````
 
 ### FAQ
 
@@ -43,22 +58,3 @@ If you encounter error about backend of librosa, try to install ffmpeg with comm
 sudo apt update
 sudo apt install ffmpeg
 ```
-
-## Example materials
-
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "/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

@@ -1,64 +0,0 @@
-# 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

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

@@ -1,95 +0,0 @@
-# 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 Prompt
+### Single Prompts
 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 import Omni
+from vllm_omni.entrypoints.omni_llm import OmniLLM
 
 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 = Omni(
+    omni_llm = OmniLLM(
         model=model_name,
         log_stats=args.enable_stats,
         log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
@@ -283,19 +283,6 @@ 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 mixed_modalities
+                  --query-type use_audio_in_video

examples/offline_inference/qwen3_omni/README.md

@@ -1,12 +1,12 @@
-# Offline Example of vLLM-Omni for Qwen3-Omni
+# Offline Example of vLLM-Omni for Qwen2.5-omni
 
 ## 🛠️ Installation
 
 Please refer to [README.md](../../../README.md)
 
-## Run examples (Qwen3-Omni)
+## Run examples (Qwen3-omni)
 ### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-Omni README.md](../qwen2_5_omni/README.md)
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-omni README.md](../qwen2_5_omni/README.md)
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -15,7 +15,7 @@ Then run the command below.
 ```bash
 bash run_multiple_prompts.sh
 ```
-### Single Prompt
+### Single Prompts
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -24,10 +24,6 @@ Then run the command below.
 ```bash
 bash run_single_prompt.sh
 ```
-If you have not enough memory, you can set thinker with tensor parallel. Just run the command below.
-```bash
-bash run_single_prompt_tp.sh
-```
 
 ### 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 import Omni
+from vllm_omni.entrypoints.omni_llm import OmniLLM
 
 SEED = 42
 
@@ -127,9 +127,8 @@ def main(args):
     model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
     query_result = query_map[args.query_type]()
 
-    omni_llm = Omni(
+    omni_llm = OmniLLM(
         model=model_name,
-        stage_configs_path=args.stage_configs_path,
     )
 
     thinker_sampling_params = SamplingParams(
@@ -277,12 +276,6 @@ def parse_args():
         default=None,
         help="Path to a .txt file with one prompt per line (preferred).",
     )
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=None,
-        help="Path to a stage configs file.",
-    )
 
     return parser.parse_args()
 
@@ -290,3 +283,7 @@ def parse_args():
 if __name__ == "__main__":
     args = parse_args()
     main(args)
+
+    # use examples:
+    # python end2end.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type text
+    # python end2end.py --model /custom_path/Qwen3-Omni-30B-A3B-Instruct --query-type use_video_only

examples/offline_inference/qwen3_omni/qwen3_omni_moe_tp.yaml

@@ -1,94 +0,0 @@
-# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
-# Stage 0: Thinker (multimodal understanding + text generation)
-# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
-# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
-
-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

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

examples/offline_inference/qwen_image/text_to_image.py

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

examples/offline_inference/tmp_test/README.md

@@ -1,34 +0,0 @@
-# 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

@@ -1,264 +0,0 @@
-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

@@ -1,45 +0,0 @@
-#!/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

@@ -1,367 +0,0 @@
-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

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

examples/offline_inference/tmp_test/ray_worker.sh

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

examples/offline_inference/tmp_test/run_multiple_prompts.sh

@@ -1,8 +0,0 @@
-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

@@ -1,11 +0,0 @@
-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

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

examples/online_serving/qwen3_omni/README.md

@@ -1,8 +1,8 @@
-# Online serving Example of vLLM-Omni for Qwen3-omni
+# Online serving Example of vLLM-omni for Qwen2.5-omni
 
 ## 🛠️ Installation
 
-Please refer to [README.md](../../../README.md)
+Please refer to [README.md](../../README.md)
 
 ## Run examples (Qwen3-Omni)
 
@@ -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/qwen3_omni
+cd examples/online_serving
 ```
 
 Send request via python
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
+python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```
 
 Send request via curl
 ```bash
-bash run_curl_multimodal_generation.sh use_image
+bash run_curl_multimodal_generation.sh mixed_modalities
 ```
 
 ### 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 AsyncOmni by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
@@ -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 AsyncOmni locally.
+- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmniLLM locally.
 - `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
 - `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
 - `--port`: Port for Gradio server (default: 7861)

examples/online_serving/qwen3_omni/gradio_demo.py

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

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

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

examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh

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

mkdocs.yml

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

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

vllm_omni/__init__.py

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

vllm_omni/core/sched/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 Request, RequestStatus, SchedulerOutput, SpecDecodingStats
+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, EngineCoreOutputs
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
 
 from vllm_omni.core.sched.output import OmniNewRequestData
 from vllm_omni.core.sched.scheduler import OmniScheduler

vllm_omni/core/sched/generation_scheduler.py

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

vllm_omni/core/sched/scheduler.py

@@ -1,15 +1,8 @@
 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):
@@ -57,177 +50,3 @@ 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.utils.platform_utils import get_diffusion_worker_class
+from vllm_omni.diffusion.worker.gpu_worker import WorkerProc
 
 logger = init_logger(__name__)
 
@@ -76,9 +76,6 @@ class DiffusionEngine:
         mp.set_start_method("spawn", force=True)
         processes = []
 
-        # Get the appropriate worker class for current device
-        worker_proc = get_diffusion_worker_class()
-
         # Launch all worker processes
         scheduler_pipe_readers = []
         scheduler_pipe_writers = []
@@ -87,7 +84,7 @@ class DiffusionEngine:
             reader, writer = mp.Pipe(duplex=False)
             scheduler_pipe_writers.append(writer)
             process = mp.Process(
-                target=worker_proc.worker_main,
+                target=WorkerProc.worker_main,
                 args=(
                     i,  # rank
                     od_config,

vllm_omni/diffusion/distributed/utils.py

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

vllm_omni/diffusion/models/__init__.py

@@ -1,3 +0,0 @@
-# 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

@@ -1,17 +0,0 @@
-# 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,32 +240,19 @@ class QwenImagePipeline(
         prefix: str = "",
     ):
         super().__init__()
-        self.od_config = od_config
         self.device = get_local_device()
         model = od_config.model
-
-        # 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
-        )
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model, subfolder="scheduler")
         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", local_files_only=local_files_only
-            )
+            self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(model, subfolder="text_encoder")
             logger.info("Loaded Qwen-Image text encoder successfully")
-            self.vae = AutoencoderKLQwenImage.from_pretrained(
-                model, subfolder="vae", local_files_only=local_files_only
-            ).to(self.device)
+            self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae").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", local_files_only=local_files_only
-            )
+            self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer")
             logger.info("Loaded Qwen-Image tokenizer successfully.")
 
         self.stage = None
@@ -405,6 +392,8 @@ 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*):
@@ -776,14 +765,10 @@ class QwenImagePipeline(
                     for name, tensor in state_dict.items():
                         yield name, tensor
 
+        model_name = "Qwen/Qwen-Image"
         try:
-            # Get model path from config or download from HF
-            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, ["*"])
-
+            # Use the arguments as verified by the user
+            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,12 +113,10 @@ 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: torch.Tensor, dim: int, theta: int = 10000):
+    def rope_params(self, index, dim, theta=10000):
         """
         Args:
-            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
+            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
         """
         assert dim % 2 == 0
         freqs = torch.outer(

vllm_omni/diffusion/worker/__init__.py

@@ -1,7 +0,0 @@
-# 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

@@ -1,7 +0,0 @@
-# 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

@@ -1,196 +0,0 @@
-# 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

@@ -1,25 +0,0 @@
-# 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

@@ -1,44 +0,0 @@
-# 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

@@ -1,145 +0,0 @@
-# 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

@@ -1,119 +0,0 @@
-# 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

@@ -1,31 +0,0 @@
-# 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

@@ -1,9 +0,0 @@
-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

@@ -1,78 +0,0 @@
-#!/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

@@ -1,332 +0,0 @@
-#!/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

@@ -1,86 +0,0 @@
-# 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

@@ -1,157 +0,0 @@
-# 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

@@ -1,314 +0,0 @@
-# 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_llm.py

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

vllm_omni/entrypoints/omni_llm.py

@@ -1,8 +1,6 @@
 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
@@ -22,11 +20,6 @@ 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
@@ -38,14 +31,9 @@ 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,
-    resolve_model_config_path,
-)
+from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -96,23 +84,14 @@ 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.config_path = resolve_model_config_path(model)
-            self.stage_configs = load_stage_configs_from_yaml(self.config_path)
+            self.stage_configs = load_stage_configs_from_model(model)
         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
 
@@ -145,140 +124,36 @@ class OmniLLM:
         self.stage_list = [st for _, st in results]
         logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
-        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._ctx = mp.get_context("spawn")
         self._stage_in_queues: list[mp.Queue] = []
         self._stage_out_queues: list[mp.Queue] = []
         self._init_sleep_seconds = max(0, int(init_sleep_seconds))
         self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
-        self._start_stages(model)
+        self._start_stage_processes(model)
         # Wait for all stages to report readiness before seeding
         self._stages_ready: set[int] = set()
         self._wait_for_stages_ready(timeout=init_timeout)
 
-    def _start_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:
+    def _start_stage_processes(self, model: str) -> None:
         for stage_id, stage in enumerate(self.stage_list):
             # Use unbounded queues to avoid deadlock when seeding many requests
-            in_q: mp.Queue = self._queue_cls()
-            out_q: mp.Queue = self._queue_cls()
+            in_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
             self._stage_in_queues.append(in_q)
             self._stage_out_queues.append(out_q)
 
             # Attach queues and start Stage-owned worker process
             stage.attach_queues(in_q, out_q)
-
-            # 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.
 
@@ -296,21 +171,10 @@ class OmniLLM:
                 )
         for stage in self.stage_list:
             try:
-                if self.worker_backend == "ray" and hasattr(stage, "_ray_actor"):
-                    import ray
-                    # Ray kill
-                    ray.kill(stage._ray_actor)
-                else:
-                    stage.stop_stage_worker()
+                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()
@@ -373,14 +237,11 @@ class OmniLLM:
         # Orchestrator keeps stage objects for input derivation
         num_stages = len(self.stage_list)
 
-        # Generate globally unique request IDs and map them to original 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)
-        }
+        # Map from request_id to original prompt
+        request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
 
         # Track per-request start time for end-to-end timing
-        _req_start_ts: dict[str, float] = {}
+        _req_start_ts: dict[int, 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)
@@ -526,90 +387,51 @@ class OmniLLM:
                         )
                         continue
                     sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
-
-                    # Check if we have a connector for this edge
-                    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:
+                        # Measure transfer size and time (encode + enqueue)
+                        size_bytes = 0
                         try:
-                            if self.worker_backend == "ray":
-                                # Ray mode: Skip SHM fallback, just send object via Queue (Ray handles serialization)
-                                # This prevents SHM access errors across nodes
-                                ipc_payload = {
-                                    "request_id": req_id,
-                                    "engine_inputs": next_inputs,
-                                    "sampling_params": sp_next,
-                                    "sent_ts": time.time(),
-                                }
-                                size_bytes = 0 # Approximation or calculate size if needed
-                                try:
-                                    # Approximate size for stats
-                                    size_bytes = len(_set(next_inputs))
-                                except:
-                                    pass
-                                t0 = time.time()
-                                # Direct submit to Ray Queue
-                            else:
-                                size_bytes = 0
-                                try:
-                                    size_bytes = len(_set(next_inputs))
-                                except Exception:
-                                    size_bytes = 0
-                                t0 = time.time()
-                                ipc_payload = _encode(
-                                    next_inputs,
-                                    getattr(self, "_shm_threshold_bytes", 65536),
-                                    obj_key="engine_inputs",
-                                    shm_key="engine_inputs_shm",
-                                )
-                                ipc_payload.update(
-                                    {
-                                        "request_id": req_id,
-                                        "sampling_params": sp_next,
-                                        "sent_ts": time.time(),
-                                    }
-                                )
-                            self.stage_list[next_stage_id].submit(ipc_payload)
-                            t1 = time.time()
-                            tx_ms = (t1 - t0) * 1000.0
-                            metrics.on_forward(
-                                stage_id,
-                                next_stage_id,
-                                req_id,
-                                int(size_bytes),
-                                float(tx_ms),
-                                bool("engine_inputs_shm" in ipc_payload),
-                            )
-                        except Exception as e:
-                            logger.warning(
-                                "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
-                                req_id,
-                                e,
-                            )
-                            self.stage_list[next_stage_id].submit(
-                                {
-                                    "request_id": req_id,
-                                    "engine_inputs": next_inputs,
-                                    "sampling_params": sp_next,
-                                }
-                            )
+                            size_bytes = len(_set(next_inputs))
+                        except Exception:
+                            size_bytes = 0
+                        t0 = time.time()
+                        ipc_payload = _encode(
+                            next_inputs,
+                            getattr(self, "_shm_threshold_bytes", 65536),
+                            obj_key="engine_inputs",
+                            shm_key="engine_inputs_shm",
+                        )
+                        ipc_payload.update(
+                            {
+                                "request_id": req_id,
+                                "sampling_params": sp_next,
+                                "sent_ts": time.time(),
+                            }
+                        )
+                        self.stage_list[next_stage_id].submit(ipc_payload)
+                        t1 = time.time()
+                        tx_ms = (t1 - t0) * 1000.0
+                        metrics.on_forward(
+                            stage_id,
+                            next_stage_id,
+                            req_id,
+                            int(size_bytes),
+                            float(tx_ms),
+                            bool("engine_inputs_shm" in ipc_payload),
+                        )
+                    except Exception as e:
+                        logger.warning(
+                            "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
+                            req_id,
+                            e,
+                        )
+                        self.stage_list[next_stage_id].submit(
+                            {
+                                "request_id": req_id,
+                                "engine_inputs": next_inputs,
+                                "sampling_params": sp_next,
+                            }
+                        )
                     logger.debug(
                         "[Orchestrator] Forwarded request %s to stage-%s",
                         req_id,
@@ -712,9 +534,8 @@ 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: Any,
+        **kwargs,
     ):
         """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_devices,
+    set_stage_gpu_devices,
 )
 from vllm_omni.inputs.data import OmniTokensPrompt
 
@@ -163,7 +163,6 @@ 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.
 
@@ -193,7 +192,6 @@ 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(
@@ -297,8 +295,6 @@ 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)
@@ -342,8 +338,6 @@ 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
 
@@ -351,7 +345,6 @@ 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:
@@ -432,10 +425,7 @@ def _stage_worker(
 
     # Device mapping
     try:
-        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)
+        set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
     except Exception as e:
         _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
 
@@ -445,18 +435,6 @@ 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})
@@ -474,9 +452,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:
@@ -513,21 +491,9 @@ def _stage_worker(
                     _in_flight_ms_by_rid[rid] = 0.0
             except Exception:
                 _in_flight_ms_by_rid[rid] = 0.0
-
-            # Resolve input data (via Connector or IPC)
-            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))
-
+            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(t, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+            _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
+            _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
             batch_request_ids.append(rid)
             if isinstance(ein, list):
                 batch_engine_inputs.extend(ein)
@@ -700,7 +666,7 @@ async def _stage_worker_async(
     import logging as _logging
     import time as _time
 
-    from vllm_omni.entrypoints.async_omni import AsyncOmniStageLLM
+    from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM
     from vllm_omni.entrypoints.log_utils import (
         compute_and_log_stage_request_stats,
         count_tokens_from_outputs,
@@ -751,10 +717,7 @@ async def _stage_worker_async(
 
     # Device mapping
     try:
-        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)
+        set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
     except Exception as e:
         _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
 

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 import AsyncOmni
+from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
 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(
+    async with build_async_omni_llm(
         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(
+async def build_async_omni_llm(
     args: Namespace,
     *,
     disable_frontend_multiprocessing: Optional[bool] = None,
     client_config: Optional[dict[str, Any]] = None,
 ) -> AsyncIterator[EngineClient]:
-    """Build an AsyncOmni instance from command-line arguments.
+    """Build an AsyncOmniLLM instance from command-line arguments.
 
-    Creates an async context manager that yields an AsyncOmni instance
+    Creates an async context manager that yields an AsyncOmniLLM instance
     configured from the provided arguments. Handles forkserver setup if
     needed and ensures proper cleanup on exit.
 
@@ -127,7 +127,7 @@ async def build_async_omni(
         client_config: Optional client configuration dictionary
 
     Yields:
-        EngineClient instance (AsyncOmni) ready for use
+        EngineClient instance (AsyncOmniLLM) 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(
         forkserver.ensure_running()
         logger.debug("Forkserver setup complete!")
 
-    # Context manager to handle async_omni lifecycle
+    # Context manager to handle async_omni_llm lifecycle
     # Ensures everything is shutdown and cleaned up on error/exit
-    async with build_async_omni_from_stage_config(
+    async with build_async_omni_llm_from_stage_config(
         args,
         disable_frontend_multiprocessing=disable_frontend_multiprocessing,
-    ) as async_omni:
-        yield async_omni
+    ) as async_omni_llm:
+        yield async_omni_llm
 
 
 @asynccontextmanager
-async def build_async_omni_from_stage_config(
+async def build_async_omni_llm_from_stage_config(
     args: Namespace,
     *,
     disable_frontend_multiprocessing: bool = False,
     client_config: Optional[dict[str, Any]] = None,
 ) -> AsyncIterator[EngineClient]:
-    """Create AsyncOmni from stage configuration.
+    """Create AsyncOmniLLM from stage configuration.
 
-    Creates an AsyncOmni instance either in-process or using multiprocess
+    Creates an AsyncOmniLLM 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_from_stage_config(
         client_config: Optional client configuration dictionary
 
     Yields:
-        EngineClient instance (AsyncOmni) ready for use
+        EngineClient instance (AsyncOmniLLM) ready for use
 
     Note:
         Stage configurations are loaded from args.stage_configs_path if provided,
@@ -182,21 +182,21 @@ async def build_async_omni_from_stage_config(
             "To disable frontend multiprocessing, set VLLM_USE_V1=0."
         )
 
-    async_omni: Optional[EngineClient] = None
+    async_omni_llm: Optional[EngineClient] = None
 
     try:
         if getattr(args, "stage_configs_path", None):
-            async_omni = AsyncOmni(model=args.model, stage_configs_path=args.stage_configs_path)
+            async_omni_llm = AsyncOmniLLM(model=args.model, stage_configs_path=args.stage_configs_path)
         else:
-            async_omni = AsyncOmni(model=args.model)
+            async_omni_llm = AsyncOmniLLM(model=args.model)
 
         # # Don't keep the dummy data in memory
         # await async_llm.reset_mm_cache()
 
-        yield async_omni
+        yield async_omni_llm
     finally:
-        if async_omni:
-            async_omni.shutdown()
+        if async_omni_llm:
+            async_omni_llm.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 (AsyncOmni)
+        engine_client: Engine client instance (AsyncOmniLLM)
         vllm_config: vLLM configuration object
         state: FastAPI application state object to initialize
         args: Parsed command-line arguments

vllm_omni/entrypoints/stage_utils.py

@@ -13,143 +13,100 @@ from omegaconf import OmegaConf
 logger = logging.getLogger(__name__)
 
 
-def set_stage_devices(
-    stage_id: int,
-    devices: str | int | None,
-    device_type: str | None = None,
-) -> None:
-    """Configure per-stage device visibility and current device (CUDA or NPU).
+def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
+    """Configure per-stage CUDA visibility and current device.
 
-    This function sets environment variables that control which devices are visible
-    to the process, and sets the current device. It must be called BEFORE worker
-    initialization so that workers see the correct devices.
-
-    Args:
-        stage_id: Stage identifier for logging
-        devices: Device specification:
-            - 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()
+    Behavior
+    - Comma-separated string (e.g. "2,5,7"): set CUDA_VISIBLE_DEVICES exactly
+      to this list; logical index 0 is used as current device.
+    - Integer or digit-string: treat as logical index (0-based) into the current
+      CUDA_VISIBLE_DEVICES mapping; map to the physical device, and then set
+      CUDA_VISIBLE_DEVICES to this single device.
+    - None/"cpu": keep default visibility.
+    - Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string.
     """
-    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[env_var] = devices
+            os.environ["CUDA_VISIBLE_DEVICES"] = devices
             toks = [t.strip() for t in devices.split(",") if t.strip() != ""]
             if toks:
                 try:
                     selected_physical = int(toks[0])
                     logger.debug(
-                        "[Stage-%s] Set %s to %s; logical 0 -> physical %s",
+                        "[Stage-%s] Set CUDA_VISIBLE_DEVICES 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 %s device: %s", stage_id, device_type_label, e)
+                    logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
                     selected_physical = None
         elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()):
             logical_idx = max(0, int(devices))
-            vis = os.environ.get(env_var)
+            vis = os.environ.get("CUDA_VISIBLE_DEVICES")
             if vis:
                 try:
                     mapping = [int(x) for x in vis.split(",") if x.strip() != ""]
                     if 0 <= logical_idx < len(mapping):
                         selected_physical = mapping[logical_idx]
                 except Exception as e:
-                    logger.debug("[Stage-%s] Failed to map logical index via %s: %s", stage_id, env_var, e)
+                    logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e)
                     selected_physical = None
             if selected_physical is None:
                 selected_physical = int(logical_idx)
-            os.environ[env_var] = str(selected_physical)
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
             logger.debug(
-                "[Stage-%s] Logical index %d -> physical %s; set %s to single device",
+                "[Stage-%s] Logical index %d -> physical %s; set CUDA_VISIBLE_DEVICES 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[env_var] = str(selected_physical)
-            logger.debug("[Stage-%s] Set %s to single device %s (fallback)", stage_id, env_var, selected_physical)
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
+            logger.debug(
+                "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical
+            )
 
         try:
-            import torch  # noqa: WPS433
+            import torch
 
-            if is_available_fn():
+            if torch.cuda.is_available():
                 try:
-                    set_device_fn(0)
+                    torch.cuda.set_device(0)
                 except Exception as e:
                     logger.debug(
-                        "[Stage-%s] %s set_device(0) failed: %s", stage_id, device_type_label, e, exc_info=True
+                        "[Stage-%s] torch.cuda.set_device(0) failed: %s",
+                        stage_id,
+                        e,
+                        exc_info=True,
                     )
-                num = device_count_fn()
+                num = torch.cuda.device_count()
                 info = []
                 for i in range(num):
-                    total = get_device_properties_fn(i).total_memory
-                    free, _ = mem_get_info_fn(i)
+                    total = torch.cuda.get_device_properties(i).total_memory
+                    free, _ = torch.cuda.mem_get_info(i)
                     info.append(
                         {
                             "idx": i,
-                            "name": get_device_name_fn(i),
+                            "name": torch.cuda.get_device_name(i),
                             "total": int(total),
                             "free": int(free),
                         }
                     )
-                logger.debug("[Stage-%s] %s devices visible=%s info=%s", stage_id, device_type_label, num, info)
+                logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info)
         except Exception as e:
-            logger.debug("[Stage-%s] Failed to query %s devices: %s", stage_id, device_type_label, e, exc_info=True)
+            logger.debug(
+                "[Stage-%s] Failed to query CUDA devices: %s",
+                stage_id,
+                e,
+                exc_info=True,
+            )
     except Exception as e:
         logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)
 

vllm_omni/entrypoints/utils.py

@@ -1,56 +1,24 @@
+from __future__ import annotations
+
+import logging
 import os
 from pathlib import Path
 
 from omegaconf import OmegaConf
 from vllm.transformers_utils.config import get_config
 
-from vllm_omni.utils import detect_device_type
+logger = logging.getLogger(__name__)
 
 # 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 and device type.
-    First tries to load a device-specific YAML file from stage_configs/{device_type}/
-    directory. If not found, falls back to the default config file.
+    Loads stage configurations based on the model type. Looks for a
+    YAML configuration file in the stage_configs directory matching
+    the model's model_type.
 
     Args:
         model: Model name or path (used to determine model_type)
@@ -61,8 +29,13 @@ def load_stage_configs_from_model(model: str) -> list:
     Raises:
         FileNotFoundError: If no stage config file exists for the model type
     """
-    stage_config_path = resolve_model_config_path(model)
-    stage_configs = load_stage_configs_from_yaml(config_path=stage_config_path)
+    hf_config = get_config(model, trust_remote_code=True)
+    model_type = hf_config.model_type
+    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
+    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,7 +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
+from vllm_omni.model_executor.preprocess_mixin import PreprocessMixin
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -50,10 +50,12 @@ logger = init_logger(__name__)
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniForConditionalGeneration(
-    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE
+    nn.Module, SupportsMultiModal, SupportsPP, SupportsMRoPE,
+    Qwen2_5OmniConditionalGenerationMixin, PreprocessMixin
 ):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
+        self.has_preprocess = False
         self.have_multimodal_outputs = True
         config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config
         multimodal_config = vllm_config.model_config.multimodal_config
@@ -84,6 +86,9 @@ class Qwen2_5OmniForConditionalGeneration(
             self.token2wav = None
 
         elif self.model_stage == "talker":
+            # register the process function for the talker stage
+            self.has_preprocess = True
+            self.set_preprocess(self.process_fn)
             self.thinker = None
             # Initialize talker model wrapper (handles projection + LM)
             self.talker = init_vllm_registered_model(
@@ -202,7 +207,6 @@ 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)
@@ -230,27 +234,15 @@ 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=thinker_input_ids,
-                positions=thinker_positions,
-                intermediate_tensors=intermediate_tensors,
-                inputs_embeds=thinker_inputs_embeds,
-                **kwargs,
-            )
+            with torch.inference_mode():
+                thinker_output = self.thinker(
+                    input_ids=input_ids,
+                    positions=positions[0],
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                    **kwargs,
+                )
 
             if isinstance(thinker_output, tuple):
                 embeds, text_hidden_states = thinker_output
@@ -265,151 +257,13 @@ class Qwen2_5OmniForConditionalGeneration(
 
         # 2) Talker (if codec not provided)
         if self.model_stage == "talker":
-            # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
-            # Rules:
-            # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
-            # - Decode segments consist of a single non-special token.
-            # - If additional_information is provided
-            # (can be a list split by request or a concatenated tensor plus a list of shapes),
-            #   then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments;
-            # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them;
-            # otherwise, keep the original embedding.
-
+            # mock data for profile
             if input_ids is None and additional_information is None:
                 input_ids = torch.zeros(
-                    inputs_embeds.shape[0],
-                    dtype=torch.long,
-                    device=inputs_embeds.device,
+                    inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device
                 )
                 additional_information = {}
                 self.thinker_reply_part = torch.zeros_like(inputs_embeds)
-                is_profile = True
-            else:
-                is_profile = False
-
-            # Ensure we have base embeddings when only ids are provided
-            if inputs_embeds is None and input_ids is not None:
-                inputs_embeds = self.talker.get_input_embeddings(input_ids)
-
-            # ------- Request-scoped additional information (no cross-request concat) -------
-            request_ids: Optional[list[str]] = kwargs.get("request_ids")  # ordered
-            request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get("request_token_spans")
-            addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id")
-            runtime_addi = kwargs.get("runtime_additional_information")
-
-            # Normalize runtime_addi into a mapping by request_id for convenience
-            runtime_addi_by_req: dict[str, dict] = {}
-            if (
-                isinstance(request_ids, list)
-                and isinstance(runtime_addi, list)
-                and len(runtime_addi) == len(request_ids)
-            ):
-                for i, rid in enumerate(request_ids):
-                    if isinstance(rid, str) and isinstance(runtime_addi[i], dict):
-                        runtime_addi_by_req[rid] = runtime_addi[i]
-            elif isinstance(request_ids, list) and isinstance(runtime_addi, dict):
-                for rid in request_ids:
-                    if isinstance(rid, str) and isinstance(runtime_addi.get(rid), dict):
-                        runtime_addi_by_req[rid] = runtime_addi[rid]
-
-            # Containers to return per-request updates (e.g., thinker_reply_part_per_request)
-            update_by_req_id: dict[str, dict] = {}
-
-            # ------- Prefill: span_len > 1 -------
-            if (
-                not is_profile
-                and isinstance(request_ids, list)
-                and isinstance(request_token_spans, list)
-                and isinstance(addi_by_req, dict)
-            ):
-                for idx_req, rid in enumerate(request_ids):
-                    s, e = request_token_spans[idx_req]
-                    span_len = int(e) - int(s)
-                    if span_len <= 1:
-                        continue
-                    info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {}
-                    if not isinstance(info, dict):
-                        info = {}
-                    pe = info.get("prompt_embeds")  # Tensor [P,H]
-                    tr = info.get("thinker_result")  # Tensor [K,H]
-                    ptoks = info.get("prompt_token_ids")  # list[int]
-                    otoks = info.get("thinker_output_token_ids")  # list[int]
-
-                    if not isinstance(pe, torch.Tensor):
-                        pe = torch.zeros(
-                            0,
-                            self.talker.config.hidden_size,
-                            dtype=inputs_embeds.dtype,
-                            device=self._module_device(self.model),
-                        )
-                    if not isinstance(tr, torch.Tensor):
-                        tr = torch.zeros(
-                            0,
-                            self.talker.config.hidden_size,
-                            dtype=inputs_embeds.dtype,
-                            device=self._module_device(self.model),
-                        )
-                    if not isinstance(ptoks, (list, torch.Tensor)):
-                        ptoks = []
-                    if not isinstance(otoks, (list, torch.Tensor)):
-                        otoks = []
-
-                    req_input_ids, req_embeds = self._thinker_to_talker_prefill(
-                        voice_type=voice_type,
-                        output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)),
-                        output_token_ids=otoks,
-                        thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)),
-                        prompt_token_ids=ptoks,
-                    )
-                    seg_len = min(span_len, req_embeds.shape[0])
-                    inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
-                    if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
-                        input_ids[s : s + seg_len] = req_input_ids
-
-                    # Prepare per-request reply queue for subsequent decode: drop first row
-                    if tr.ndim == 2 and tr.shape[0] > 0:
-                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = (
-                            tr[1:].detach().to("cpu").contiguous()
-                        )
-
-            # ------- Decode: span_len == 1 -------
-            if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list):
-                for idx_req, rid in enumerate(request_ids):
-                    s, e = request_token_spans[idx_req]
-                    if (int(e) - int(s)) != 1:
-                        continue
-                    # choose step vector in priority order
-                    step_vec = None
-                    # A) runtime queue
-                    q = None
-                    if isinstance(rid, str):
-                        q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request")
-                    if isinstance(q, torch.Tensor) and q.numel() > 0:
-                        step_vec = q[0:1]
-                        new_q = q[1:].detach().to("cpu").contiguous()
-                        update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q
-                    else:
-                        # B) per-request provided decode vector (optional)
-                        info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {}
-                        dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None
-                        if isinstance(dv, torch.Tensor) and dv.numel() > 0:
-                            step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
-                        elif (
-                            hasattr(self, "thinker_reply_part")
-                            and isinstance(self.thinker_reply_part, torch.Tensor)
-                            and self.thinker_reply_part.numel() > 0
-                        ):
-                            # C) fallback shared pool
-                            step_vec = self.thinker_reply_part[0:1]
-                            self.thinker_reply_part = self.thinker_reply_part[1:]
-
-                    if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
-                        one_id = input_ids[s : s + 1]
-                        _, one_embed = self._thinker_to_talker_decode_one_step(
-                            output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
-                            output_token_ids=one_id,
-                        )
-                        inputs_embeds[s] = one_embed[0]
 
             with torch.inference_mode():
                 talker_hidden = self.talker(
@@ -417,10 +271,6 @@ class Qwen2_5OmniForConditionalGeneration(
                     positions=positions[0],
                     inputs_embeds=inputs_embeds,
                 )
-            multimodal_outputs: dict = None
-            # Return updates if any
-            if update_by_req_id:
-                multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id}
 
             if sampling_metadata is not None:
                 # the padding token id is set to text model's pad token id,
@@ -429,7 +279,7 @@ class Qwen2_5OmniForConditionalGeneration(
 
             return OmniOutput(
                 text_hidden_states=talker_hidden,
-                multimodal_outputs=multimodal_outputs,
+                multimodal_outputs=None,
             )
 
         if self.model_stage == "code2wav":
@@ -741,6 +591,68 @@ class Qwen2_5OmniForConditionalGeneration(
             return talker_hf_config.tts_text_start_token_id
         return self.tts_text_spk_token_ids[voice_type]
 
+    def process_fn(self,
+        input_ids: torch.Tensor,
+        inputs_embeds: torch.Tensor,
+        **info_dict: object,
+    ):
+        # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
+        # Rules:
+        # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
+        # - Decode segments consist of a single non-special token.
+        # - If additional_information is provided (can be a list split by request or a 
+        #   concatenated tensor plus a list of shapes), then for each request, reconstruct 
+        #   the thinker→talker input embeddings for the Prefill segments;
+        # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional),
+        #   add them; otherwise, keep the original embedding.
+
+        # Ensure we have base embeddings when only ids are provided
+        if inputs_embeds is None and input_ids is not None:
+            inputs_embeds = self.talker.get_input_embeddings(input_ids)
+
+        span_len = input_ids.shape[0]
+        if span_len > 1:
+            # prefill
+            return self.thinker_to_talker_process(input_ids, inputs_embeds, **info_dict)
+        else:
+            # decode
+            return self.thinker_to_talker_decode_one_step(input_ids, inputs_embeds, **info_dict)
+
+
+    def thinker_to_talker_process(self, 
+        input_ids: torch.Tensor,
+        inputs_embeds: torch.Tensor,
+        **info_dict: object,
+    ):
+        update_dict = {}
+
+        prompt_embeds = info_dict.get("prompt_embeds")  # Tensor [P,H]
+        thinker_result = info_dict.get("thinker_result")  # Tensor [K,H]
+        prompt_token_ids = info_dict.get("prompt_token_ids")  # list[int]
+        thinker_output_token_ids = info_dict.get("thinker_output_token_ids")  # list[int]
+
+        if not isinstance(prompt_embeds, torch.Tensor):
+            prompt_embeds = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
+        if not isinstance(thinker_result, torch.Tensor):
+            thinker_result = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model))
+        if not isinstance(prompt_token_ids, (list, torch.Tensor)):
+            prompt_token_ids = []
+        if not isinstance(thinker_output_token_ids, (list, torch.Tensor)):
+            thinker_output_token_ids = []
+        
+        req_input_ids, req_embeds = self._thinker_to_talker_prefill(
+            voice_type=info_dict.get("voice_type"),
+            output_prompt_embeds=thinker_result.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+            output_token_ids=thinker_output_token_ids,
+            thinker_prompt_embeds=prompt_embeds.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+            prompt_token_ids=prompt_token_ids,
+        )
+        
+        if thinker_result.ndim == 2 and thinker_result.shape[0] > 0:
+            update_dict["thinker_reply_part"] = thinker_result[1:].detach().to("cpu").contiguous()
+
+        return req_input_ids, req_embeds, update_dict
+
     def _thinker_to_talker_prefill(
         self,
         voice_type: str,
@@ -786,6 +698,38 @@ class Qwen2_5OmniForConditionalGeneration(
             )
         return prompt_token_ids_processed, prompt_embeds
 
+    def thinker_to_talker_decode_one_step(self,
+        input_ids,
+        inputs_embeds,
+        **info_dict
+    ):
+        update_dict = {}
+        # choose step vector in priority order
+        step_vec = None
+        q = info_dict.get("thinker_reply_part", None)
+        if isinstance(q, torch.Tensor) and q.numel() > 0:
+            step_vec = q[0:1]
+            new_q = q[1:].detach().to("cpu").contiguous()
+            update_dict["thinker_reply_part"] = new_q
+        else:
+            # B) per-request provided decode vector (optional)
+            dv = info_dict.get("decode_output_prompt_embeds") if isinstance(info_dict, dict) else None
+            if isinstance(dv, torch.Tensor) and dv.numel() > 0:
+                step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
+            elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0:
+                # C) fallback shared pool
+                step_vec = self.thinker_reply_part[0:1]
+                self.thinker_reply_part = self.thinker_reply_part[1:]
+
+        if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
+            one_id = input_ids[0:1]
+            _, one_embed = self._thinker_to_talker_decode_one_step(
+                output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
+                output_token_ids=one_id,
+            )
+            inputs_embeds[0] = one_embed[0]
+        return input_ids[0:1], inputs_embeds[0:1], update_dict
+
     def _thinker_to_talker_decode_one_step(
         self,
         output_prompt_embeds,
@@ -866,7 +810,7 @@ class Qwen2_5OmniForConditionalGeneration(
         __init__."""
         if self.token2wav is None or self.token2wav_config is None:
             return
-        device = self._module_device(self.token2wav)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         # optional speaker resources
         conds = getattr(self.token2wav_config, "conds", None)
         ref_mels = getattr(self.token2wav_config, "ref_mels", None)

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py

@@ -33,8 +33,6 @@ 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):
@@ -726,13 +724,7 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
     else:
         beta = 0.0
 
-    # 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)
+    kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
 
     # Compute time indices
     if is_even:
@@ -742,7 +734,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)
+        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)  # Ensures correct shape
 
     sinc_filter = torch.sinc(2 * cutoff * time_indices)
     normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
@@ -753,29 +745,6 @@ 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__()
@@ -791,28 +760,14 @@ class UpSample1d(nn.Module):
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
-        if is_npu():
-            # TODO: When F.pad supports replicate mode on NPU, remove this branch
-            input_dtype = hidden_states.dtype
-            # F.pad in NPU doesn't support BF16 when mode is replicate.
-            # To ensure the accuracy, manually pad the input tensor.
-            hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad, self.pad)
-            filter_convert_dtype = self.filter.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_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
@@ -838,29 +793,14 @@ class DownSample1d(nn.Module):
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
-        if is_npu():
-            input_dtype = hidden_states.dtype
-            # F.pad in NPU doesn't support BF16 when mode is replicate.
-            # To ensure the accuracy, manually pad the input tensor.
-            hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad_left, self.pad_right)
-            filter_on_device = self.filter.to(device=hidden_states.device, dtype=hidden_states.dtype)
-            out = F.conv1d(
-                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)
+        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 onto the current span
+                    # req_embeds: [S,D], req_input_ids: [S]  →  map到当前 span
                     seg_len = min(span_len, req_embeds.shape[0])
-                    # Force-align dtype/device/memory layout to avoid mismatches from precision or layout differences
+                    # 强制对齐 dtype/device/内存布局，避免精度或布局差异导致比较不相等
                     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,7 +6,6 @@ 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
@@ -165,48 +164,21 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
 
         from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 
-        # Try attention backends in order of preference, with runtime error handling
-        # This handles cases where the backend is registered but not actually available
-        attention_backends = ["flash_attention_2", "xformers", "eager", "sdpa"]
-        attn_output = None
-        last_error = None
-
-        for backend_name in attention_backends:
-            if backend_name not in ALL_ATTENTION_FUNCTIONS:
-                continue
-
-            try:
-                attention_interface = ALL_ATTENTION_FUNCTIONS[backend_name]
-                attn_output, _ = attention_interface(
-                    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."
-            )
+        attention_interface = ALL_ATTENTION_FUNCTIONS["flash_attention_2"]
+        attn_output, _ = attention_interface(
+            self,
+            q_heads,
+            k_heads,
+            v_heads,
+            None,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.head_dim**-0.5,
+            sliding_window=None,
+            use_cache=use_cache,
+            position_ids=position_ids,
+            output_hidden_states=True,
+            output_attentions=False,
+        )
         attn_output = attn_output.reshape(*(hidden_states.shape[:-1]), -1).contiguous()
 
         attn_output = self.o_proj(attn_output)
@@ -401,22 +373,16 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
         past_key_values: Optional[Any] = None,
         use_cache: Optional[bool] = False,
         cache_position: Optional[torch.LongTensor] = None,
-        **kwargs: Any,
-    ) -> Any:
+        **kwargs,
+    ):
         """
         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:
-            Named tuple with .last_hidden_state and .past_key_values attributes
+            Object with .last_hidden_state attribute
         """
         batch_size, seq_len, _ = inputs_embeds.shape
 
@@ -455,6 +421,8 @@ 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,5 +1,4 @@
 from collections.abc import Iterable
-from typing import Any
 
 import torch
 import torch.nn as nn
@@ -285,7 +284,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
 
         return result_codes, summed_embeddings
 
-    def init_multi_modal(self, thinker_config: Any) -> None:
+    def init_multi_modal(self, thinker_config):
         """
         Initialize multimodal components from the thinker.
 

vllm_omni/model_executor/preprocess_mixin.py

@@ -0,0 +1,27 @@
+from collections.abc import Callable
+import torch
+
+class PreprocessMixin:
+    """
+    Mixin class for all stages in the Omni model.
+    """
+    def set_preprocess(self, preprocess_fn: Callable) -> None:
+        """
+        Set a preprocess function for the stage.
+        Args:
+            preprocess_fn: The preprocess function to register.
+        """
+        self.preprocess = preprocess_fn
+
+    def preprocess(self,
+                input_ids: torch.Tensor,
+                inputs_embeds: torch.Tensor,
+                **input_dict: object
+    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+        """
+        Process the input_ids and inputs_embeds for the given input_dict.
+        Returns the processed input_ids, inputs_embeds, and the input_dict.
+
+        If the stage don't applicable, return the original input_ids, inputs_embeds, and an empty dict.
+        """
+        return input_ids, inputs_embeds, {}

vllm_omni/model_executor/stage_configs/npu/qwen2_5_omni.yaml

@@ -1,94 +0,0 @@
-# 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,9 +26,6 @@ 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
@@ -55,22 +52,17 @@ 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: "2"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
-      gpu_memory_utilization: 0.3
+      gpu_memory_utilization: 0.15
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false
@@ -86,9 +78,6 @@ 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:
@@ -96,31 +85,6 @@ 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 Any, Union
+from typing import Union
 
 import torch
 from vllm.inputs import TextPrompt
@@ -47,11 +47,11 @@ def _compute_talker_prompt_ids_length(info):
 
 
 def thinker2talker(
-    stage_list: list[Any],
-    engine_input_source: list[int],
-    prompt: Union[OmniTokensPrompt, TextPrompt, None] = None,
+    stage_list,
+    engine_input_source,
+    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
     requires_multimodal_data: bool = False,
-) -> list[OmniTokensPrompt]:
+):
     """
     Process thinker outputs to create talker inputs.
 
@@ -64,7 +64,6 @@ 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
@@ -112,11 +111,11 @@ def thinker2talker(
 
 
 def talker2code2wav(
-    stage_list: list[Any],
-    engine_input_source: list[int],
-    prompt: Union[OmniTokensPrompt, TextPrompt, None] = None,
+    stage_list,
+    engine_input_source,
+    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
     requires_multimodal_data: bool = False,
-) -> list[OmniTokensPrompt]:
+):
     """
     Process talker outputs to create code2wav inputs.
 
@@ -129,7 +128,6 @@ 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

@@ -1,11 +0,0 @@
-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

@@ -1,54 +0,0 @@
-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_ar_model_runner.py

@@ -356,4 +356,4 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 merged[k] = v
         setattr(req_state, "additional_information_cpu", merged)
 
-    # ===== Helper functions extracted for clarity and reuse =====
+    # ===== Helper functions extracted for clarity and reuse =====

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,8 +6,6 @@ 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
@@ -35,63 +33,6 @@ 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.
@@ -698,6 +639,34 @@ class OmniGPUModelRunner(GPUModelRunner):
             encoder_inputs = self._extract_encoder_inputs(scheduler_output)
             model_kwargs.update(encoder_inputs)
 
+        if self.model.has_preprocess:
+            # 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.get(req_id)
+                req_infos = (getattr(req_state, "additional_information_cpu", None)
+                             if req_state is not None else None)
+
+                start_offset = int(self.query_start_loc.cpu[req_index])
+                sched_tokens = int(num_scheduled_tokens_np[req_index])
+                s, e = start_offset, start_offset + sched_tokens
+                span_len = int(e) - int(s)
+
+                # call the custom process function
+                req_input_ids, req_embeds, update_dict = self.model.preprocess(
+                    input_ids=input_ids[s:e],
+                    inputs_embeds=inputs_embeds[s:e],
+                    **req_infos
+                )
+                # TODO(Peiqi): the merge stage could move out from the critical path
+                self._merge_additional_information_update(req_id, update_dict)
+
+                # update the inputs_embeds and input_ids
+                seg_len = min(span_len, req_embeds.shape[0])
+                inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
+                if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
+                    input_ids[s : s + seg_len] = req_input_ids
+
         return (
             num_scheduled_tokens,
             num_input_tokens,

vllm_omni/worker/npu/npu_ar_model_runner.py

@@ -1,963 +0,0 @@
-# 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

@@ -1,15 +0,0 @@
-# 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

@@ -1,748 +0,0 @@
-# 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

@@ -1,15 +0,0 @@
-# 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

@@ -1,543 +0,0 @@
-# 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