hsliu
51d2570464
- Add comprehensive PRD, architecture design, and test design documents - Implement core modules: OmniLLM, AsyncOmniLLM, stage configurations - Add DiT scheduler and cache manager for diffusion models - Implement CLI integration with --omni flag support - Add API server and plugin system for vLLM integration - Create comprehensive test suite with fixtures - Update dependencies to vLLM 0.10.2 and PyTorch 2.8.0 - Add conda environment setup and package installation - Implement stage-based processing architecture - Add multimodal output processing capabilities This commit establishes the foundation for multi-modality models inference and serving with non-autoregressive structures.
180 lines
6.2 KiB
Markdown
180 lines
6.2 KiB
Markdown
# 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
|