vllm-omni/docs/PRD.md

# vLLM-omni Product Requirements Document (PRD)

## 1. Product Overview

### 1.1 Product Name
vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures

### 1.2 Product Vision
Extend vLLM beyond traditional text-based, autoregressive generation to support multi-modality models with non-autoregressive structures and non-textual outputs while maintaining vLLM's proven architecture and performance.

### 1.3 Target Users
- AI researchers working with multimodal models
- ML engineers building production inference systems
- Developers integrating DiT (Diffusion Transformer) models
- Organizations requiring efficient multimodal model serving

## 2. Core Requirements

### 2.1 Functional Requirements

#### 2.1.1 Multi-Stage Processing
- **REQ-001**: Support stage-based model processing where each stage can use different engine types (AR/DiT)
- **REQ-002**: Enable sequential processing through multiple stages with data flow between stages
- **REQ-003**: Support both autoregressive (AR) and diffusion (DiT) model stages

#### 2.1.2 vLLM Compatibility
- **REQ-004**: Maintain full compatibility with vLLM V1 architecture (AsyncLLM and EngineCore patterns)
- **REQ-005**: Support existing vLLM CLI commands with `--omni` flag extension
- **REQ-006**: Reuse vLLM's multiprocess worker architecture for scalability

#### 2.1.3 Multimodal Support
- **REQ-007**: Support text, image, and latent space inputs/outputs
- **REQ-008**: Enable image-to-image, text-to-image, and text-to-text generation
- **REQ-009**: Support hidden state passing between AR and DiT stages

#### 2.1.4 CLI and API
- **REQ-010**: Provide CLI command: `vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8000`
- **REQ-011**: Support both online (AsyncOmniLLM) and offline (OmniLLM) inference modes
- **REQ-012**: Maintain vLLM's existing API compatibility

### 2.2 Non-Functional Requirements

#### 2.2.1 Performance
- **REQ-013**: Maintain vLLM's inference performance for AR stages
- **REQ-014**: Optimize DiT stage processing with caching mechanisms
- **REQ-015**: Support distributed inference across multiple GPUs

#### 2.2.2 Scalability
- **REQ-016**: Support horizontal scaling through vLLM's worker process pattern
- **REQ-017**: Enable efficient memory management for large multimodal models
- **REQ-018**: Support batch processing for multiple requests

#### 2.2.3 Extensibility
- **REQ-019**: Easy integration of new modalities and model architectures
- **REQ-020**: Pluggable scheduler and executor components
- **REQ-021**: Support for future non-autoregressive model types

## 3. Technical Architecture

### 3.1 Core Components

#### 3.1.1 Entry Points
- **OmniServeCommand**: CLI wrapper that intercepts vLLM commands with `--omni` flag
- **OmniLLM**: Offline inference class supporting multi-stage processing
- **AsyncOmniLLM**: Online inference class with asynchronous processing

#### 3.1.2 Stage Management
- **OmniStageConfig**: Configuration for each processing stage
- **Stage Engine List**: Multiple AsyncLLM instances for each stage
- **Stage I/O Management**: Data flow between stages

#### 3.1.3 Engine Components
- **EngineCore**: Reused from vLLM (no changes needed)
- **OmniDiffusionScheduler**: New scheduler for DiT models
- **DiTCacheManager**: Caching system for DiT optimization
- **MultiprocExecutor**: Reused from vLLM for DiT without diffusers
- **DiffusersPipelineExecutor**: New executor for diffusers integration

#### 3.1.4 Model Runners
- **OmniDiffusionModelRunner**: Handles DiT model execution
- **OmniARModelRunner**: Handles AR model execution with hidden state output
- **ModelRunnerOutput**: Extended to support multimodal outputs via pooler_output

#### 3.1.5 Output Processing
- **MultimodalOutputProcessor**: Handles final multimodal output processing
- **RequestState**: Extended to support pooling outputs
- **Output Handlers**: Type-specific output processing

### 3.2 Data Flow
```
API Server → OmniLLM/AsyncOmniLLM → LLMEngine/AsyncLLM → Engine Core
→ Scheduler (AR/DiT) → Executor (AR/DiT) → Worker (AR/DiT)
→ ModelRunner (AR/DiT) → RequestState → OutputProcessor → Final Output
```

## 4. Implementation Phases

### Phase 1: Foundation (Weeks 1-2)
- Package structure and dependencies
- Basic OmniLLM and AsyncOmniLLM classes
- Stage configuration system
- CLI integration

### Phase 2: Core Components (Weeks 3-4)
- DiT scheduler implementation
- Model runners for AR and DiT
- Basic output processing

### Phase 3: Advanced Features (Weeks 5-6)
- Caching system implementation
- Multimodal output processing
- Request state management

### Phase 4: Integration & Testing (Weeks 7-8)
- End-to-end integration
- Comprehensive testing
- Performance optimization
- Documentation

## 5. Success Criteria

### 5.1 Functional Success
- [ ] Successfully run `vllm serve model --omni` command
- [ ] Process multi-stage AR→DiT pipelines
- [ ] Generate multimodal outputs (text + image)
- [ ] Maintain vLLM API compatibility

### 5.2 Performance Success
- [ ] AR stage performance within 5% of native vLLM
- [ ] DiT stage processing with reasonable latency
- [ ] Memory usage comparable to vLLM for equivalent models

### 5.3 Quality Success
- [ ] 90%+ test coverage
- [ ] All integration tests passing
- [ ] Documentation complete and accurate

## 6. Risk Assessment

### 6.1 Technical Risks
- **High**: vLLM API changes breaking compatibility
- **Medium**: DiT model integration complexity
- **Low**: Performance overhead from multi-stage processing

### 6.2 Mitigation Strategies
- Regular vLLM compatibility testing
- Incremental DiT integration with fallback options
- Performance benchmarking at each stage

## 7. Dependencies

### 7.1 External Dependencies
- vLLM >= 0.10.2
- PyTorch >= 2.7
- Transformers >= 4.30.0
- FastAPI, Uvicorn for API serving
- Ray for distributed computing

### 7.2 Optional Dependencies
- xDiT for DiT acceleration
- Cache-DiT for advanced caching
- Diffusers for pipeline-based DiT models

## 8. Future Roadmap

### 8.1 Short-term (3 months)
- Additional DiT model support
- Performance optimizations
- Enhanced caching strategies

### 8.2 Medium-term (6 months)
- Support for video generation models
- Advanced scheduling strategies
- Multi-GPU DiT optimization

### 8.3 Long-term (12 months)
- Custom model architecture support
- Advanced multimodal fusion
- Production deployment tools