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