From 81f04d3cf9f796c7cbbb986a15c63c80c7f29840 Mon Sep 17 00:00:00 2001 From: hsliu Date: Sun, 28 Sep 2025 21:58:48 +0800 Subject: [PATCH] update ar_dit_test --- .gitignore | 1 - docs/architecture/detailed arch design.md | 488 ++++-------------- docs/architecture/high level arch design.md | 2 +- .../implementation_architecture.md | 55 +- docs/implementation_summary.md | 293 ----------- docs/testing/test_design.md | 35 +- examples/README.md | 55 +- examples/basic/README.md | 78 +-- examples/basic/docker_example.py | 178 ------- .../{simple_usage.py => text_generation.py} | 0 examples/omni/README.md | 31 ++ examples/omni/ar_dit_diffusers.py | 175 +++++++ examples/omni/configs/ar_dit_local.yaml | 28 + pyproject.toml | 4 +- requirements.txt | 1 + scripts/docker-setup-macos.sh | 83 --- scripts/download_models.py | 147 ------ scripts/quick_test.sh | 82 --- scripts/test_serving.sh | 86 ++- vllm_omni/config/stage_config.py | 87 ++-- vllm_omni/core/sched/__init__.py | 1 - vllm_omni/core/sched/diffusion_scheduler.py | 229 -------- vllm_omni/core/stage_manager.py | 5 +- vllm_omni/engine/diffusion_engine.py | 132 +++++ vllm_omni/engine/output_processor.py | 78 ++- vllm_omni/entrypoints/omni.py | 17 +- vllm_omni/entrypoints/omni_llm.py | 310 ++++++++--- vllm_omni/executor/diffusers_executor.py | 64 +++ vllm_omni/model_executor/__init__.py | 0 vllm_omni/plugin.py | 166 ------ .../worker/gpu_diffusion_model_runner.py | 133 +++++ vllm_omni/worker/gpu_diffusion_worker.py | 56 ++ 32 files changed, 1192 insertions(+), 1908 deletions(-) delete mode 100644 docs/implementation_summary.md delete mode 100644 examples/basic/docker_example.py rename examples/basic/{simple_usage.py => text_generation.py} (100%) create mode 100644 examples/omni/README.md create mode 100644 examples/omni/ar_dit_diffusers.py create mode 100644 examples/omni/configs/ar_dit_local.yaml delete mode 100755 scripts/docker-setup-macos.sh delete mode 100755 scripts/download_models.py delete mode 100755 scripts/quick_test.sh delete mode 100644 vllm_omni/core/sched/diffusion_scheduler.py create mode 100644 vllm_omni/engine/diffusion_engine.py create mode 100644 vllm_omni/executor/diffusers_executor.py create mode 100644 vllm_omni/model_executor/__init__.py delete mode 100644 vllm_omni/plugin.py create mode 100644 vllm_omni/worker/gpu_diffusion_model_runner.py create mode 100644 vllm_omni/worker/gpu_diffusion_worker.py diff --git a/.gitignore b/.gitignore index 94b0c7b0c..5941d027c 100644 --- a/.gitignore +++ b/.gitignore @@ -238,7 +238,6 @@ Dockerfile.dev # Kubernetes k8s/ -*.yaml *.yml !configs/*.yaml !configs/*.yml \ No newline at end of file diff --git a/docs/architecture/detailed arch design.md b/docs/architecture/detailed arch design.md index 268aba8a2..20c759ba1 100644 --- a/docs/architecture/detailed arch design.md +++ b/docs/architecture/detailed arch design.md @@ -2,21 +2,22 @@ ## Overview -vLLM-omni is a multi-modality extension for vLLM that supports non-autoregressive structures and non-textual outputs. This document outlines the key software abstractions, APIs, and dependencies for the system, designed to maximize reuse of vLLM's proven architecture. +vLLM-omni extends vLLM with multi-stage, multimodal capabilities. It supports both autoregressive (AR) stages and diffusion (DiT) stages, enabling pipelines such as text→image generation. ## Architecture Principles -1. **vLLM V1 Compatibility**: Built on vLLM's Engine V1 architecture with AsyncLLM and EngineCore patterns -2. **Stage-based Processing**: Models are divided into multiple stages, each processed by different Engine Cores -3. **Multiple Engine Core Support**: Each stage can use either AR Engine Core (reusing vLLM) or Diffusion Engine Core (new DiT support), or other new Cores -4. **Worker Process Pattern**: Follows vLLM's multiprocess worker architecture for scalability -5. **Extensibility**: Easy integration of new modalities, model architectures, and output formats +1. **vLLM V1 Compatibility** – Reuse LLMEngine/AsyncLLM and executor abstractions. +2. **Stage-based Processing** – Each stage encapsulates its own engine and configuration. +3. **Executor/Worker Separation** – Executors remain thin; model-specific logic lives in workers and runners. +4. **Extensibility** – New modalities can be added by implementing dedicated stage configs, workers, and output processors. ## Key Data Flow -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 --> OutputProcessor (New one for final multimodal output) +``` +Request → OmniLLM/AsyncOmniLLM → StageManager → LLMEngine/Executor → Worker → ModelRunner → OutputProcessor → Response +``` + +For diffusion stages, the executor hands off to `DiffusionGPUWorker`, which wraps `DiffusionModelRunner` (responsible for diffusers pipeline loading and inference). ## Core Components (Ordered by Data Flow) @@ -27,399 +28,130 @@ pip install vllm-omni ``` ### 2. Online inference launch (Entry Point) - -Keep consistency with the vllm main branch ```bash -vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8000 +vllm serve Qwen/Qwen3-0.6B --omni --port 8000 ``` -#### Design Architecture -```mermaid -graph TD - A[vllm serve model --omni] --> B[vLLM-omni CLI Wrapper] - B --> C{Detect --omni flag} - C -->|Yes| D[Parse omniConfig] - C -->|No| E[Forward to vLLM CLI] - D --> F[Initialize AsyncOmniLLM] - F --> G[Start Omni Server] - G --> H[Multi-stage Processing] - E --> I[Standard vLLM Pipeline] - - subgraph "vLLM-omni Components" - F - G - H - end - - subgraph "vLLM Components" - I - end -``` - -#### Omni Serve Command Implementation -File: vllm_omni/entrypoints/cli/main.py +### 3. Omni Serve Command (CLI) ```python -import sys -import argparse -from typing import List, Optional -from vllm_omni.entrypoints.omni import OmniServeCommand - -def main(): - """Main CLI entry point that intercepts vLLM commands""" - # Check if --omni flag is present - if "--omni" in sys.argv: - # Remove --omni flag and process with vLLM-omni - omni_args = [arg for arg in sys.argv[1:] if arg != "--omni"] - omni_serve = OmniServeCommand() - omni_serve.run(omni_args) - else: - # Forward to original vLLM CLI - from vllm.entrypoints.cli.main import main as vllm_main - vllm_main() - -if __name__ == "__main__": - main() +# vllm_omni/entrypoints/cli/main.py +if "--omni" in sys.argv: + omni_args = [arg for arg in sys.argv[1:] if arg not in {"--omni", "serve"}] + OmniServeCommand().run(omni_args) +else: + from vllm.entrypoints.cli.main import main as vllm_main + vllm_main() ``` -#### vLLM Plugin Integration -File: vllm_omni/pyproject.toml +### 4. Offline Inference (OmniLLM) ```python -[project.scripts] -# Override vLLM's CLI entry point -vllm = "vllm_omni.cli:main" - -# Add entry point for vLLM integration -[project.entry-points."vllm.plugins"] -omni = "vllm_omni.plugin:OmniPlugin" -``` - -### 3. Offline inference launch (Entry Point) -Design an upper level class to incorporate multi Engines, each engine has a engine call. -```python -from vllm.entrypoints.llm import LLM -from vllm.v1.engine.llm_engine import LLMEngine - +# vllm_omni/entrypoints/omni_llm.py class OmniLLM(LLM): - """Extended LLM supporting multiple engines and stage-based processing""" - - def __init__(self, stage_configs: List[OmniStageConfig]): - super().__init__() - self.stage_configs = stage_configs - self.engine_list = [] # List of AsyncLLM instances for each stage + def __init__(..., stage_configs): + super().__init__(model=stage_configs[0].model_path if stage_configs else "test-model") + self.stage_manager = StageManager(stage_configs) 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: - stage_llm = LLMEngine.from_vllm_config( - vllm_config=stage_config.vllm_config, - executor_class=self.executor_class, - log_stats=self.log_stats - ) - self.engine_list.append(stage_llm) - - def generate( - self, - stage_args_list: List[stage_args], - use_tqdm: Union[bool, Callable[..., tqdm]] = True, - lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None, - priority: Optional[list[int]] = None, - ) -> list[RequestOutput]: - """Main generation interface - orchestrates multi-stage processing""" + self.stage_manager.initialize_engines() - # Process through each stage sequentially - for i, stage_config in enumerate(self.stage_configs): - stage_engine = self.engine_list[i] - - # Prepare input for this stage - stage_args = stage_args_list[i] - prompt_str, engine_request, tokenization_kwargs = self._process_stage_inputs(stage_config, **stage_args) - - # Add inputs to Engine - stage_engine.add_request(requesy_id, prompt_str, tokenization_kwargs) - # Run Engine - stage_output = stage_engine.step() - - # Update stage input and output of each stage for later usage - self._update_stage_io(stage_output, stage_config) - - - # Process final output - final_output = self.output_processor.process_output( - self.stage_configs[-1].stage_output - ) - return final_output - - def _process_stage_inputs(stage_config, **stage_args) -> Any: - """Prepare input for specific stage""" - if stage_config.engine_type == "AR": - return self._process_ar_inputs(**stage_args) - elif stage_config.engine_type == "DiT": - return self._process_dit_inputs(**stage_args) - else: - raise NotImplementedError - - def _process_dit_inputs(**stage_args)-> Any: - image_latent = self.vae.encode(**stage_args) - - return image_latent + def generate(...): + current_output = None + for idx, (config, args) in enumerate(zip(self.stage_configs, stage_args_list)): + engine = self.stage_manager.get_engine(idx) + processed = self._process_stage_inputs(config, args, current_output) + stage_output = self._execute_stage(engine, processed, stage_config=config) + current_output = stage_output + config.stage_output = stage_output + return self.output_processor.process_output(current_output) ``` -### 4. Stage config for setting of different model stages +### 5. Stage Manager ```python +# vllm_omni/core/stage_manager.py +class StageManager: + def initialize_engines(self): + if self._initialized: + return + for config in self.stage_configs: + if config.engine_type == "AR": + engine = self._create_ar_engine(config) + elif config.engine_type == "DiT": + engine = self._create_dit_engine(config) + self.engine_list.append(engine) + self._initialized = True +``` + +### 6. Diffusion Execution (DiT Stage) +```python +# vllm_omni/worker/diffusion_model_runner.py +class DiffusionModelRunner: + def __init__(self, model_path, pipeline_name=None, device=None, dtype=None): + self.device, self.torch_dtype = self._resolve_device_and_dtype(device, dtype) + self._pipeline = self._load_pipeline(model_path, pipeline_name) + + def generate(...): + output = self._pipeline(...) + return DiffusionRunnerOutput(prompt=prompt, images=getattr(output, "images", []) or []) +``` + +```python +# vllm_omni/worker/gpu_diffusion_worker.py +class DiffusionGPUWorker: + def __init__(self, model_path, ...): + self.model_runner = DiffusionModelRunner(model_path, ...) + + def generate(...): + return self.model_runner.generate(...) +``` + +```python +# vllm_omni/executor/diffusers_executor.py +class DiffusersPipelineExecutor: + def __init__(self, model_path, ...): + self.worker = DiffusionGPUWorker(model_path, ...) + + def generate(...): + return self.worker.generate(...) +``` + +This mirrors the standard executor→worker→runner layering in vLLM. + +### 7. Stage Configuration +```python +# vllm_omni/config/stage_config.py @dataclass class OmniStageConfig: - """Configuration for a processing stage""" stage_id: int - engine_type: str # "AR" or "DiT" + engine_type: Literal["AR", "DiT"] model_path: str input_modalities: List[str] output_modalities: List[str] - vllm_config: Optional[VllmConfig] = None # For engine config of corresponding stage - executor_class: type[Executor] # For execute class config of corresponding stage - dit_config: Optional[DiTConfig] = None # For diffusion stages - cache_config: Optional[DiTCacheConfig] = None + vllm_config: Optional[VllmConfig] = None + executor_class: Optional[type[Executor]] = None + dit_config: Optional[DiTConfig] = None ``` -For AR stage, the setting is: -```python -from vllm.v1.core.sched.scheduler import Scheduler -from vllm.v1.executor.multiproc_executor import MultiprocExecutor -ar_stage_config = OmniStageConfig() -ar_stage_config.vllm_config.scheduler_config.scheduler_cls = Scheduler # original vllm scheduler for AR -ar_stage_config.executor_class = MultiprocExecutor # original vllm executor for AR -``` -For DiT stage, the setting is: -```python -from vllm.v1.executor.multiproc_executor import MultiprocExecutor -dit_stage_config = OmniStageConfig() -dit_stage_config.vllm_config.scheduler_config.scheduler_cls = DiffusionScheduler # New scheduler for DiT -# For diffusion models without using diffusers -dit_stage_config.executor_class = MultiprocExecutor +- AR stages default to `UniProcExecutor`. +- DiT stages set `executor_class` to `DiffusersPipelineExecutor` when `use_diffusers=True`. -# For diffusion models using diffusers -dit_stage_config.executor_class = DiffusersPipelineExecutor -``` - -### 5. Online inference main class - -Similar to OmniLLM in offline inference, add some asynchronous processing, referring to AsyncLLM -```python -from vllm.v1.engine.async_llm import AsyncLLM - - -class AsyncOmniLLM(AsyncLLM): - """Extended AsyncLLM supporting multiple engines and stage-based processing""" - - def __init__(self, stage_configs: List[OmniStageConfig]): - super().__init__() -``` - -### 6. Engine Core -No need to change. The specific executor, scheduler and output type can be transferred into it with new configs. -```python -# class: EngineCore -# EngineCore.step (simplified) -class EngineCore: - def step(self): - scheduler_output = scheduler.schedule() - model_output = executor.execute_model(scheduler_output) - engine_outputs = scheduler.update_from_output( - scheduler_output, model_output - ) - return engine_outputs -``` - -### 7. Scheduler -Create a new Diffusion Scheduler for DiT, which is inherited from original vllm AR scheduler. At first, no complicated strategy. -Just process the request first in first out. Then create a class of DiT Cache Manager for optimization in the future. -```python -from vllm.v1.core.sched.scheduler import SchedulerInterface - - -class OmniDiffusionScheduler(SchedulerInterface): - - def __init__( - self, - vllm_config: VllmConfig, - kv_cache_config: KVCacheConfig, - dit_cache_config: DiTCacheConfig, - structured_output_manager: StructuredOutputManager, - mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, - include_finished_set: bool = False, - log_stats: bool = False, - ) -> None: - - self.dit_cache_manager = DiTCacheManager(dit_cache_config) -``` -For DiT Cache Manager (Can refer to xDiT): -```python -class DiTCacheManager: - """Manages DiT-specific caching""" - - def __init__(self, config: DiTCacheConfig): - self.cache_tensors: Dict[str, torch.Tensor] = {} - self.cache_groups: List[DiTCacheTensor] = config.dit_cache_tensors - - def allocate_cache(self, request_id: str, size: int) -> torch.Tensor - def get_cache(self, request_id: str) -> Optional[torch.Tensor] - def release_cache(self, request_id: str) -> None - def clear_expired_cache(self) -> None -``` - -### 8. Executor for DiT without diffusers - -No need to change. Just import original AR executor. - - - -### 9. Worker - -#### Inheritance strategy -Prefer reusing the mature GPU Worker end-to-end. Worker class is selected by configuration (vllm_config). Do not add a new executor-specific worker binding. If customization is needed, override only `init_device` to construct the `OmniDiffusionModelRunner`; all other behaviors (device init details, profiling, sleep/wake, PP/TP comms, execute path) remain from `vllm/v1/worker/gpu_worker.py::Worker`. - - -Inherited and overridden -```python -# Optional: only if you need to plug a custom OmniDiffusionModelRunner. -from vllm.v1.worker.gpu_worker import Worker as GPUWorker -from vllm.v1.worker.diffusion_model_runner import OmniDiffusionModelRunner - -class Worker(GPUWorker): - def init_device(self) -> None: - #those related to device check and init - ... - - self.model_runner = OmniDiffusionModelRunner(self.vllm_config, ...) -``` - -### 10. Model Runner - -### Function map (Model Runner) -#### 1) Inherited and overridden -Those parts relied to the KV Cache will be omitted if we do not register the model to the vllm config. The engine core will view it as do not require KV Cache, and handle it properly - -Reuse `vllm/v1/outputs.py::ModelRunnerOutput`： -- OmniDiffusionModelRunner: Use the `pooler_output=[Tensor,...]` to return multi modal tensors -- OmniARModelRunner: Use the `pooler_output=[Tensor,...]` to return hidden states. -```python -from typing import Optional, Union -import torch -from vllm.v1.worker.gpu_model_runner import GPUModelRunner -from vllm.v1.outputs import ModelRunnerOutput - - -class OmniDiffusionModelRunner(GPUModelRunner): - @torch.inference_mode() - def execute_model( - self, - scheduler_output: "SchedulerOutput", - intermediate_tensors: Optional[IntermediateTensors] = None, - ) -> Union[ModelRunnerOutput, IntermediateTensors]: - ... - return ModelRunnerOutput( - req_ids=[...], - req_id_to_index={...}, - sampled_token_ids=[], - spec_token_ids=None, - logprobs=None, - prompt_logprobs_dict={}, - pooler_output=[Tensor,...], # Return Hidden states - kv_connector_output=None, - num_nans_in_logits=None, - )# return multi modal tensors via pooler_output=[Tensor,...] - - -class OmniARModelRunner(GPUModelRunner): - @torch.inference_mode() - def execute_model( - self, - scheduler_output: "SchedulerOutput", - intermediate_tensors: Optional[IntermediateTensors] = None, - ) -> Union[ModelRunnerOutput, IntermediateTensors]: - ... - return ModelRunnerOutput( - req_ids=[...], - req_id_to_index={...}, - sampled_token_ids=[], - spec_token_ids=None, - logprobs=None, - prompt_logprobs_dict={}, - pooler_output=[Tensor,...], # Return Hidden states - kv_connector_output=None, - num_nans_in_logits=None, - )# return hidden states via pooler_output=[Tensor,...] -``` - -### 11. RequestState Role and Usage -```bash -# RequestState serves as the per-request state tracker in OutputProcessor: -# - Maintains request-specific state (tokens, logprobs, detokenizer, etc.) -# - Converts EngineCoreOutput → RequestOutput/PoolingRequestOutput -# - Manages request lifecycle from registration to completion -AsyncLLM.add_request() -↓ -OutputProcessor.add_request() -↓ -RequestState.from_new_request() → 创建请求状态 - -AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() → 创建一个Background loop 持续从EngineCore获取输出 -↓ -OutputProcessor.process_outputs() → 更新状态并处理输出 -↓ -RequestState.make_request_output() → 转换为最终输出,格式为RequestOutput或者PoolingRequestOutput -↓ -RequestOutputCollector.put() → 推送到队列（AsyncLLM） -``` - -Need to add implementation for an existing method -```python -class RequestState: - def __init__(self, request_id: str, parent_req: Optional[ParentRequest], - request_index: int, lora_name: Optional[str], - output_kind: RequestOutputKind, prompt: Optional[str], - prompt_token_ids: list[int], logprobs_processor: Optional[LogprobsProcessor], - detokenizer: Optional[IncrementalDetokenizer], - max_tokens_param: Optional[int], arrival_time: float, - queue: Optional[RequestOutputCollector], log_stats: bool): - - def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput: - """Create PoolingOutput for multimodal/pooling requests""" - return PoolingOutput(data=pooling_output) -``` - - -### 12. OutputProcessor - -For hidden state output, original OutputProcessor already support. - -We only need to add new one for final multimodal output. +### 8. Output Processing ```python +# vllm_omni/engine/output_processor.py class MultimodalOutputProcessor(OutputProcessor): - """Handles multimodal output processing""" - - def __init__(self): - self.output_handlers: Dict[str, OutputProcessor] = {} - - def process_outputs(self, engine_core_outputs: list[EngineCoreOutput], ...): - for engine_core_output in engine_core_outputs: - # Option 1: Use output_type field (if EngineCoreOutput is extended) - if engine_core_output.output_type == "image": - self._process_image_output(engine_core_output) - elif engine_core_output.output_type == "text+image": - self._process_text_image_output(engine_core_output) - elif engine_core_output.output_type == "latents": - self._process_latents_output(engine_core_output) - elif engine_core_output.output_type == "text": - self._process_text_output(engine_core_output) - else: - # Fallback: use existing pooling_output logic - if engine_core_output.pooling_output is not None: - self._process_pooling_output(engine_core_output) - else: - self._process_text_output(engine_core_output) + def process_output(self, engine_core_output): + output_type = self._detect_output_type(engine_core_output) + handler = self.output_handlers.get(output_type, self._process_pooling_output) + return handler(engine_core_output) ``` +### 9. CLI Package Configuration +```toml +# pyproject.toml +[project.scripts] +vllm = "vllm_omni.entrypoints.cli.main:main" +vllm-omni = "vllm_omni.entrypoints.cli.main:main" +``` - +## Future Work +- Scheduler adaptations for diffusion batching. +- Advanced output processors for audio/video. +- Multi-worker execution (ray / multiprocess) for heavy DiT stages. diff --git a/docs/architecture/high level arch design.md b/docs/architecture/high level arch design.md index 0306606c0..1545ca877 100644 --- a/docs/architecture/high level arch design.md +++ b/docs/architecture/high level arch design.md @@ -136,4 +136,4 @@ Hongsheng Liu: [hsliuustc@gmail.com](mailto:hsliuustc@gmail.com) [image10]: -[image11]: \ No newline at end of file +[image11]: diff --git a/docs/architecture/implementation_architecture.md b/docs/architecture/implementation_architecture.md index 60af881e0..9a1c4a82d 100644 --- a/docs/architecture/implementation_architecture.md +++ b/docs/architecture/implementation_architecture.md @@ -12,13 +12,12 @@ vllm_omni/ │ └── cache_config.py # Caching configurations ├── core/ # Core processing components │ ├── __init__.py -│ ├── omni_llm.py # OmniLLM and AsyncOmniLLM │ ├── stage_manager.py # Multi-stage orchestration │ ├── dit_cache_manager.py # DiT caching system │ └── sched/ # Schedulers │ ├── __init__.py │ ├── scheduler.py # Base scheduler interface -│ └── diffusion_scheduler.py # DiT scheduler + ├── engine/ # Engine components │ ├── __init__.py │ ├── processor.py # Output processing @@ -29,14 +28,14 @@ vllm_omni/ │ └── diffusers_executor.py # Diffusers pipeline executor ├── model_executor/ # Model execution components │ ├── __init__.py -│ ├── ar_model_runner.py # OmniARModelRunner -│ └── dit_model_runner.py # OmniDiffusionModelRunner ├── worker/ # Worker implementations │ ├── __init__.py -│ └── omni_worker.py # Extended worker for DiT +│ ├── 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 @@ -87,24 +86,29 @@ class OmniStageConfig: input_modalities: List[str] output_modalities: List[str] vllm_config: Optional[VllmConfig] = None - executor_class: type[Executor] = MultiprocExecutor dit_config: Optional[DiTConfig] = None - cache_config: Optional[DiTCacheConfig] = None + executor_class: Optional[type[Executor]] = None stage_output: Optional[Any] = None ``` ### 3.2 DiT Configuration ```python -# vllm_omni/config/dit_config.py +# vllm_omni/config/stage_config.py @dataclass class DiTConfig: - model_type: str - scheduler_type: str + 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 + guidance_scale: float = 7.5 use_diffusers: bool = False diffusers_pipeline: Optional[str] = None + height: int = 512 + width: int = 512 ``` ### 3.3 Cache Configuration @@ -179,34 +183,9 @@ class AsyncOmniLLM(AsyncLLM): """Async generation interface""" # Similar to OmniLLM but with async/await patterns pass -``` +`` -### 4.3 DiT Scheduler Implementation - -```python -# vllm_omni/core/sched/diffusion_scheduler.py -class OmniDiffusionScheduler(SchedulerInterface): - def __init__( - self, - vllm_config: VllmConfig, - kv_cache_config: KVCacheConfig, - dit_cache_config: DiTCacheConfig, - structured_output_manager: StructuredOutputManager, - mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, - include_finished_set: bool = False, - log_stats: bool = False, - ) -> None: - super().__init__(vllm_config, kv_cache_config, structured_output_manager, - mm_registry, include_finished_set, log_stats) - self.dit_cache_manager = DiTCacheManager(dit_cache_config) - - def schedule(self) -> SchedulerOutput: - """Schedule DiT requests with caching optimization""" - # Implement DiT-specific scheduling logic - pass -``` - -### 4.4 Model Runners +### 4.3 Model Runners ```python # vllm_omni/model_executor/dit_model_runner.py diff --git a/docs/implementation_summary.md b/docs/implementation_summary.md deleted file mode 100644 index cdd3b4b45..000000000 --- a/docs/implementation_summary.md +++ /dev/null @@ -1,293 +0,0 @@ -# vLLM-omni Implementation Summary - -## Overview - -This document provides a comprehensive summary of the vLLM-omni implementation plan and current progress. The implementation follows a structured approach with PRD, architecture design, test design, and code implementation phases. - -## Implementation Status - -### ✅ Completed Phases - -#### 1. Product Requirements Document (PRD) -- **File**: `docs/PRD.md` -- **Status**: Complete -- **Contents**: - - Product vision and target users - - Functional and non-functional requirements - - Technical architecture overview - - Implementation phases and success criteria - - Risk assessment and mitigation strategies - -#### 2. Architecture Design -- **File**: `docs/architecture/implementation_architecture.md` -- **Status**: Complete -- **Contents**: - - Detailed package structure - - Core module dependencies - - Configuration system design - - Implementation details for all components - - Memory management and error handling strategies - -#### 3. Test Design -- **File**: `docs/testing/test_design.md` -- **Status**: Complete -- **Contents**: - - Comprehensive testing strategy - - Unit, integration, and E2E test specifications - - Performance and compatibility testing - - Test configuration and execution guidelines - -#### 4. Package Setup -- **Status**: Complete -- **Components**: - - Updated `pyproject.toml` with vLLM integration - - Package structure and dependencies - - CLI entry point configuration - - vLLM plugin system setup - -#### 5. Core Modules -- **Status**: Complete -- **Components**: - - `vllm_omni/config/`: Configuration management - - `vllm_omni/core/omni_llm.py`: OmniLLM and AsyncOmniLLM classes - - `vllm_omni/core/stage_manager.py`: Multi-stage orchestration - - `vllm_omni/core/dit_cache_manager.py`: DiT caching system - - `vllm_omni/core/sched/diffusion_scheduler.py`: DiT scheduler - - `vllm_omni/engine/output_processor.py`: Multimodal output processing - -### 🚧 In Progress / Pending - -#### 6. Scheduler and Executor Components -- **Status**: Partially Complete -- **Completed**: - - DiT scheduler implementation - - DiT cache manager -- **Pending**: - - Executor implementations - - Integration with vLLM's executor system - -#### 7. Model Runners -- **Status**: Pending -- **Components**: - - `OmniDiffusionModelRunner`: DiT model execution - - `OmniARModelRunner`: AR model execution with hidden states - - Integration with vLLM's model runner system - -#### 8. Output Processing -- **Status**: Partially Complete -- **Completed**: - - Basic multimodal output processor -- **Pending**: - - RequestState extensions - - Advanced output handling - -#### 9. CLI Integration -- **Status**: Complete -- **Components**: - - CLI entry point with `--omni` flag support - - OmniServeCommand implementation - - API server for serving - - vLLM plugin system - -#### 10. Testing and Validation -- **Status**: Partially Complete -- **Completed**: - - Basic test structure - - Configuration tests -- **Pending**: - - Integration tests - - E2E tests - - Performance benchmarks - -## Key Implementation Details - -### Package Structure -``` -vllm_omni/ -├── config/ # Configuration management -│ ├── stage_config.py # OmniStageConfig, DiTConfig, etc. -│ └── __init__.py -├── core/ # Core processing components -│ ├── omni_llm.py # OmniLLM and AsyncOmniLLM -│ ├── stage_manager.py # Multi-stage orchestration -│ ├── dit_cache_manager.py # DiT caching system -│ └── sched/ # Schedulers -│ └── diffusion_scheduler.py -├── engine/ # Engine components -│ └── output_processor.py # Multimodal output handling -├── entrypoints/ # Entry points and CLI -│ ├── cli/main.py # CLI main entry point -│ ├── omni.py # OmniServeCommand -│ └── api_server.py # API server -├── plugin.py # vLLM plugin system -└── __init__.py # Main package exports -``` - -### Core Classes - -#### OmniLLM -- **Purpose**: Offline multi-stage processing -- **Features**: - - Sequential stage processing - - AR and DiT stage support - - Input/output processing between stages - - Integration with vLLM's LLMEngine - -#### AsyncOmniLLM -- **Purpose**: Online multi-stage processing -- **Features**: - - Asynchronous stage processing - - Real-time request handling - - Integration with vLLM's AsyncLLM - -#### StageManager -- **Purpose**: Orchestrate multiple stage engines -- **Features**: - - Engine creation and management - - Stage configuration handling - - Resource cleanup - -#### DiTCacheManager -- **Purpose**: Optimize DiT inference with caching -- **Features**: - - Tensor caching with multiple strategies (FIFO, LRU, LFU) - - Memory management - - Cache statistics and monitoring - -### Configuration System - -#### OmniStageConfig -- **Purpose**: Configure individual processing stages -- **Features**: - - Support for AR and DiT engine types - - Modality specification - - Executor class selection - - DiT-specific configuration - -#### DiTConfig -- **Purpose**: Configure DiT-specific parameters -- **Features**: - - Model type and scheduler configuration - - Inference parameters (steps, guidance scale) - - Diffusers integration support - -### CLI Integration - -#### Command Line Interface -```bash -# Basic usage -vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8000 - -# Multi-stage configuration -vllm serve model --omni --ar-stage text-model --dit-stage image-model - -# DiT-specific options -vllm serve model --omni --dit-steps 50 --dit-guidance-scale 7.5 -``` - -#### API Server -- **Endpoints**: - - `POST /generate`: Generate text or multimodal content - - `GET /health`: Health check - - `GET /info`: Service information -- **Features**: - - FastAPI-based REST API - - Async request handling - - Multi-stage processing support - -## Installation and Usage - -### Installation -```bash -# Install vLLM first -pip install vllm>=0.2.0 - -# Install vLLM-omni -pip install vllm-omni - -# Or install from source -git clone https://github.com/hsliuustc0106/vllm-omni -cd vllm-omni -pip install -e ".[dev]" -``` - -### Basic Usage -```python -from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config - -# Create stage configurations -ar_config = create_ar_stage_config( - stage_id=0, - model_path="Qwen/Qwen3-0.6B", - input_modalities=["text"], - output_modalities=["text"] -) - -dit_config = create_dit_stage_config( - stage_id=1, - model_path="stabilityai/stable-diffusion-2-1", - input_modalities=["text"], - output_modalities=["image"] -) - -# Create OmniLLM instance -omni_llm = OmniLLM([ar_config, dit_config]) - -# Generate with multi-stage processing -stage_args = [ - {"prompt": "A beautiful landscape"}, - {"prompt": "A beautiful landscape"} -] - -outputs = omni_llm.generate(stage_args) -``` - -## Next Steps - -### Immediate Priorities -1. **Complete Model Runners**: Implement OmniDiffusionModelRunner and OmniARModelRunner -2. **Executor Integration**: Complete executor implementations and vLLM integration -3. **Testing**: Implement comprehensive test suite -4. **Documentation**: Create user guides and API documentation - -### Medium-term Goals -1. **Performance Optimization**: Optimize memory usage and inference speed -2. **Advanced Features**: Implement advanced scheduling and caching strategies -3. **Model Support**: Add support for more model architectures -4. **Production Readiness**: Add monitoring, logging, and deployment tools - -### Long-term Vision -1. **Multi-GPU Support**: Optimize for distributed inference -2. **Custom Architectures**: Support for custom model architectures -3. **Advanced Multimodal**: Enhanced multimodal fusion capabilities -4. **Ecosystem Integration**: Integration with popular ML frameworks - -## Technical Considerations - -### vLLM Compatibility -- Maintains compatibility with vLLM V1 architecture -- Reuses proven components (scheduler, executor, worker patterns) -- Minimal modifications to existing vLLM codebase - -### Memory Management -- Efficient caching system for DiT models -- Memory sharing between stages -- Garbage collection optimization - -### Extensibility -- Plugin-based architecture -- Easy integration of new modalities -- Configurable stage processing - -### Performance -- Optimized for both AR and DiT models -- Efficient multi-stage processing -- Scalable to multiple GPUs - -## Conclusion - -The vLLM-omni implementation provides a solid foundation for multi-modality model inference with non-autoregressive structures. The current implementation covers the core functionality needed for basic multi-stage processing, with a clear path forward for advanced features and optimizations. - -The modular architecture ensures maintainability and extensibility, while the vLLM integration provides a proven foundation for production deployment. The comprehensive testing strategy and documentation ensure reliability and ease of use. - -This implementation represents a significant step forward in making advanced multimodal AI models accessible and efficient for production use cases. diff --git a/docs/testing/test_design.md b/docs/testing/test_design.md index 7710dc0e0..fdc69db79 100644 --- a/docs/testing/test_design.md +++ b/docs/testing/test_design.md @@ -100,9 +100,9 @@ class TestOmniStageConfig: def test_dit_stage_config_creation(self): """Test DiT stage configuration creation""" dit_config = DiTConfig( - model_type="dit", - scheduler_type="ddpm", - num_inference_steps=50 + num_inference_steps=50, + guidance_scale=7.5, + use_diffusers=False, ) config = OmniStageConfig( stage_id=1, @@ -130,27 +130,24 @@ class TestOmniStageConfig: #### test_dit_config.py ```python import pytest -from vllm_omni.config.dit_config import DiTConfig +from vllm.config import ModelConfig +from vllm_omni.config import DiTConfig class TestDiTConfig: def test_dit_config_defaults(self): """Test DiT configuration with defaults""" - config = DiTConfig( - model_type="dit", - scheduler_type="ddpm", - num_inference_steps=50 - ) + config = DiTConfig(num_inference_steps=50) + assert config.model_config is None assert config.use_diffusers is False assert config.diffusers_pipeline is None def test_diffusers_config(self): """Test DiT configuration with diffusers""" config = DiTConfig( - model_type="dit", - scheduler_type="ddpm", + model_config=ModelConfig(model="stabilityai/stable-diffusion-2-1"), num_inference_steps=50, use_diffusers=True, - diffusers_pipeline="stable-diffusion" + diffusers_pipeline="stable-diffusion", ) assert config.use_diffusers is True assert config.diffusers_pipeline == "stable-diffusion" @@ -431,6 +428,7 @@ class TestCLIIntegration: #### test_full_pipeline.py ```python import pytest +from vllm.config import ModelConfig from vllm_omni.core.omni_llm import OmniLLM from vllm_omni.config.stage_config import OmniStageConfig, DiTConfig @@ -454,9 +452,10 @@ class TestFullPipeline: input_modalities=["text"], output_modalities=["image"], dit_config=DiTConfig( - model_type="dit", - scheduler_type="ddpm", - num_inference_steps=10 # Reduced for testing + model_config=ModelConfig(model="stabilityai/stable-diffusion-2-1"), + num_inference_steps=10, + guidance_scale=7.5, + use_diffusers=False, ) ) ] @@ -631,9 +630,9 @@ def sample_stage_configs(): input_modalities=["text"], output_modalities=["image"], dit_config=DiTConfig( - model_type="dit", - scheduler_type="ddpm", - num_inference_steps=10 + model_config=ModelConfig(model="test-dit-model"), + num_inference_steps=10, + guidance_scale=7.5, ) ) ] diff --git a/examples/README.md b/examples/README.md index f09a1749c..2f5c91616 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1,49 +1,16 @@ # vLLM-omni Examples -This directory contains examples demonstrating how to use vLLM-omni for multi-modality model inference and serving. +A curated set of runnable samples demonstrating vLLM-omni usage. -## Directory Structure +## Directory Overview -- `basic/` - Simple usage examples and API client -- `advanced/` - Advanced features and configurations -- `multimodal/` - Multi-modality specific examples +- **basic/** – Quick-start scripts for text-only flows. + - `text_generation.py` – Minimal CPU-friendly text generation with vLLM. + - `api_client.py` – Calls the FastAPI server using the JSON API. + - `docker_run.py` – Template for launching the server in Docker. +- **omni/** – Multi-stage AR → DiT pipelines powered by diffusers. + - `ar_dit_diffusers.py` – Runs Qwen3 (AR) feeding Stable Diffusion 2.1 (DiT). + - `configs/ar_dit_local.yaml` – Sample plugin config referencing local weights. +- **advanced/** – Placeholder for upcoming complex flows (Ray, benchmarking). -## Quick Start - -1. **Basic Usage:** - ```bash - cd examples/basic - python simple_usage.py - ``` - -2. **API Server:** - ```bash - # Start server - vllm serve Qwen/Qwen3-0.6B --omni --port 8000 - - # Test with client - python api_client.py - ``` - -## Examples Overview - -### Basic Examples (`basic/`) -- `simple_usage.py` - Direct library usage (sync/async) -- `api_client.py` - HTTP API client examples -- `README.md` - Basic examples documentation - -### Advanced Examples (`advanced/`) -- Multi-stage processing -- Custom configurations -- Performance optimization - -### Multi-modality Examples (`multimodal/`) -- Text-to-image generation -- Image-to-text processing -- Multi-modal pipelines - -## Requirements - -- Python 3.8+ -- vLLM-omni installed -- For API examples: HTTP server running \ No newline at end of file +See the subdirectory READMEs for usage instructions. diff --git a/examples/basic/README.md b/examples/basic/README.md index 6a0257fbb..4e4dbb849 100644 --- a/examples/basic/README.md +++ b/examples/basic/README.md @@ -1,76 +1,10 @@ -# Basic vLLM-omni Examples +# Basic Examples -This directory contains simple examples demonstrating how to use vLLM-omni for text generation. +Quick-start scripts for experimenting with vLLM-omni. -## Examples +## Scripts -### 1. `simple_usage.py` - Direct Library Usage - -Demonstrates how to use vLLM-omni directly in your Python code: - -```bash -# Run the example -python simple_usage.py -``` - -**Features:** -- Synchronous text generation -- Asynchronous text generation -- Multi-prompt processing -- Shows both text output and token information - -### 2. `api_client.py` - API Server Usage - -Demonstrates how to interact with the vLLM-omni API server: - -```bash -# First, start the server -vllm serve Qwen/Qwen3-0.6B --omni --port 8000 - -# Then run the client (in another terminal) -python api_client.py -``` - -**Features:** -- Health check -- Text generation via HTTP API -- Multiple prompt types -- Full response structure inspection - -## Quick Start - -1. **Install vLLM-omni:** - ```bash - pip install -e . - ``` - -2. **Run direct usage example:** - ```bash - cd examples/basic - python simple_usage.py - ``` - -3. **Run API server example:** - ```bash - # Terminal 1: Start server - vllm serve Qwen/Qwen3-0.6B --omni --port 8000 - - # Terminal 2: Run client - cd examples/basic - python api_client.py - ``` - -## Expected Output - -The examples will show: -- ✅ Generated text responses -- ✅ Token counts and completion status -- ✅ Stage processing information -- ✅ Error handling - -## Requirements - -- Python 3.8+ -- vLLM-omni installed -- For API examples: HTTP server running +- `text_generation.py` – Minimal standalone usage showing how to create an LLM and generate text locally. +- `api_client.py` – Demonstrates calling the FastAPI `/generate` endpoint from Python. +Run any script with `python examples/basic/