diff --git a/docs/.nav.yml b/docs/.nav.yml
index e8281148b..fd77a1a9b 100644
--- a/docs/.nav.yml
+++ b/docs/.nav.yml
@@ -31,7 +31,6 @@ nav:
- Design Documents:
- design/index.md
- design/architecture_overview.md
- - design/vllm_omni_design.md
- Feature Design:
- design/feature/multi_request_streaming_design.md
- design/feature/omni_connector_design.md
diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md
index 3a219a564..36b524854 100644
--- a/docs/design/architecture_overview.md
+++ b/docs/design/architecture_overview.md
@@ -1,113 +1,172 @@
# Architecture Overview
-# Introduction
+This document outlines the architectural design for vLLM-Omni.
-This document outlines the architectural design for vLLM-Omni, a new repository developed under the vLLM project. vLLM-Omni supports multi-modality models with non-autoregressive structures and non-textual outputs, extending vLLM beyond traditional text-based, autoregressive generation.
+
+
+
# Goals
-The primary goals for vLLM-Omni are:
+The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support:
-* **Support for Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, alongside text.
-* **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.
-* **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.
+* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, alongside text.
+* **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.
+* **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.
* **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.
-# Expected Gains
-* It breaks through the limitation of vLLM that only supports autoregressive text generation models, and provides architectural support for popular model structures, especially multimodal generation.
-* By supporting new output formats and new model architectures, developers can easily connect to mature ecosystems such as Higress Gateway and Ray, improving the ease of deployment and efficiency of corresponding models.
-* Integrating efficient cache management, multiple parallel strategies, disaggregation from vLLM, to improve inference efficiency of corresponding models.
-
-# Non-Autoregressive Structure
-
-* **Diffusion Transformer:** initialized with a random latent noise, the transformer model iteratively predicts latent output for multiple steps, under control of sampling timestamps. Finally, it is decoded by the decoder of Variational AutoEncoder (VAE). It is not supported by vLLM now.
-* **Autoregressive:** dominant generation paradigm for text, now supported by vLLM. It generates tokens conditioned on previous tokens one by one. The efficient KV Cache management provided by vLLM can accelerate inference effectively.
- ![][image1] ![][image2]
-
-
-| | Autoregressive | DiT |
-| :---- | ----- | ----- |
-| Scenarios | Text generation | Multimodal generation |
-| Generation | Token-by-token, left to right | Continuous vector, iteratively refine |
-| Key advantages | Smooth, good local consistency | Global consistency, diversity and controlability |
-| Efficiency | More efficient with KV cache | Need generation repeatedly, slow |
-| Error accumulation | Easy to accumulate | Can fix possibly |
-| Key challenges | Limited global planning, long context generation | Large global computational cost |
-| Sequence length | Varied length | Fixed length |
-| Attention mask | Lower triangular | Full bi-directional matrix |
-| Parallel strategy | TP, DP, PP, .etc | DP, CP, .etc |
-| vLLM support | Supported | Unsupported |
-
-# Targeted models to support
+# Representative omni-modality models
According to analysis for current popular open-source models, most of them have the combination of AR+DiT. Specifically, they can be further categorized into 3 types below:
-* **DiT as a main structure, with AR as text encoder (e.g.: Qwen-Image)**
- A popular model for image generation and editing. Main visual generation models are similar to the structure, such as Flux.
- ![][image3]
-* **AR as a main structure, with DiT as multi-modal generator (e.g. BAGEL)**
- A unified multimodal comprehension and generation model. Visual generation can utilize CoT text generation output. Main unified multimodal models are similar to the structure.
- ![][image4]
-* **Extension, multiple AR+DiT (e.g. Qwen-omni)**
- A new model for multimodal input and output. It designs the Thinker-Talker-Codec structure, which is a double-AR+DiT format.
- ![][image5]
+**DiT as a main structure, with AR as text encoder (e.g.: Qwen-Image)**
+ A powerful image generation foundation model capable of complex text rendering and precise image editing.
+
+
+
+**AR as a main structure, with DiT as multi-modal generator (e.g. BAGEL)**
+ A unified multimodal comprehension and generation model, with cot text output and visual generation.
+
+
+
+**AR+DiT (e.g. Qwen-Omni)**
+ A natively end-to-end omni-modal LLM for multimodal inputs (text/image/audio/video...) and outputs (text/audio...).
+
+
+
-# Architectural Overview
+# vLLM-Omni main architecture
-vLLM-Omni will import core modules from existing vLLM framework and introduce several new components and modifications to achieve its goals. The core architecture will be based on a modular design, allowing for independent development and integration of different modality handlers and output generators.
-![][image6]
-
-## General ideas
-
-* It divides models in vLLM-Omni into multiple stages, which are processed by different Engine Cores. Each stage can choose AR Engine Core for AR model part, and Diffusion Engine Core for diffusion model part.
-* Designed a specialized diffusion engine core for DiT or dLLM.
-* Plug-in with external DiT acceleration modules such as xDiT, Cache-DiT.
-* Output linked with Higress API, Gradio for online demo, ComfyUI for less-code process.
+
+
+
## Key Components
-* **Processor:** Input processing, tokenization or encoding
-* **Scheduler:** for order and queue of incoming requests
-* **Executor:** Prepare model input and run
-* **Multimodal output processor:** support multimodal outputs
-* **Hidden State Output processor:** support intermediate hidden vectors as request output
-* **Engine Core:** accelerate AR modules, imported from vLLM main repo
-* **Diffusion Engine Core:** Wrap and accelerate DiT modules.
-* **Diffusion model:** Wrap up diffusion models with vLLM format
+| Component | Description |
+|-----------|-------------|
+| **OmniRouter** | provide an intelligent router for Omni-modality requests dispatch |
+| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
+| **AR** | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management |
+| **Diffusion** | natively implemented and optimized using acceleration components |
+| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages |
-# Examples data flow of specific models
+Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
-1. Qwen-omni
- ![][image7]![][image8]
-2. Qwen-Image
- ![][image9] ![][image10]
+## Main features
-# Overview of function call dependency
+vLLM-Omni aims to be fast, flexible, and easy to use with the following features:
-![][image11]
+### Performance and Acceleration
+The framework achieves high performance through several optimization techniques:
-[image1]:
+* **Efficient AR Support:** Leverages efficient KV cache management inherited from vLLM.
+* **Pipelined Execution:** Uses pipelined stage execution overlapping to ensure high throughput.
+* **Full Disaggregation:** Relies on the OmniConnector and dynamic resource allocation across stages.
+* **Diffusion Acceleration:** Includes integrated support for diffusion acceleration. This is managed by the acceleration layer, which handles:
+ * **Cache:** Includes DBCache, TeaCache and third-party integration(e.g., [cache-dit](https://github.com/vipshop/cache-dit)).
+ * **Parallelism:** Supports TP, CP, USP, and CFG.
+ * **Attention:** Provides an interface for third-party integration (e.g., FA3, SAGE, MindIE-SD).
+ * **Quantization:** Supports various quantization implementations including FP8 and AWQ.
+ * **FusedOps:** Allows for custom and third-party integration.
-[image2]:
+### Flexibility and Usability
-[image3]:
+vLLM-Omni is designed to be flexible and straightforward for users:
-[image4]:
+* **Heterogeneous Pipeline Abstraction:** Manages complex model workflows effectively.
+* **Hugging Face Integration:** Offers seamless integration with popular Hugging Face models.
+* **Distributed Inference:** Supports tensor, pipeline, data, and expert parallelism.
+* **Streaming Outputs:** Supports streaming outputs.
+* **Unified API:** Provides a consistent and unified API interface compatible with vLLM.
+* **OpenAI-compatible API Server:** Includes a FastAPI-based server for online serving that is compatible with the OpenAI API.
-[image5]:
+# Interface design
-[image6]:
+If you use vLLM, then you know how to use vLLM-Omni from Day 0:
-[image7]:
+
-[image8]:
+Taking **Qwen3-Omni** as an example:
-[image9]:
+## Offline Inference
+The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
-[image10]:
+```
+# Create an omni_lm with HF model name.
+from vllm_omni.entrypoints.omni import Omni
-[image11]:
+omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
+
+# Example prompts.
+om_inputs = {"prompt": prompt,
+ "multi_modal_data": {
+ "video": video_frames,
+ "audio": audio_signal,
+ }}
+
+# Generate texts and audio from the multi-modality inputs.
+outputs = omni_lm.generate(om_inputs, sampling_params_list)
+```
+
+## Online Serving
+Similar to vLLM, vLLM-Omni also provides a FastAPI-based server for online serving. Users can launch the server using the vllm serve command with the `--omni` flag:
+
+```
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
+Users can send requests to the server using curl:
+
+```
+# prepare user content
+user_content='[
+ {
+ "type": "video_url",
+ "video_url": {
+ "url": "'"$SAMPLE_VIDEO_URL"'"
+ }
+ },
+ {
+ "type": "text",
+ "text": "Why is this video funny?"
+ }
+ ]'
+ sampling_params_list='[
+ '"$thinker_sampling_params"',
+ '"$talker_sampling_params"',
+ '"$code2wav_sampling_params"'
+ ]'
+ mm_processor_kwargs="{}"
+
+# send the request
+curl -sS -X POST http://localhost:8091/v1/chat/completions \
+ -H "Content-Type: application/json" \
+ -d @- < OmniLLM/AsyncOmni (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)
-
-## Core Components (Ordered by Data Flow)
-
-### 1. Installation
-```bash
-pip install vllm
-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
-```
-
-#### 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 AsyncOmni]
- 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
-```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 Plugin Integration
-File: vllm_omni/pyproject.toml
-```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
-
-class OmniLLM(LLM):
- """Extended LLM supporting multiple engines and stage-based processing"""
-
- def __init__(self, stage_configs: List[StageConfig]):
- super().__init__()
- self.stage_configs = stage_configs
- self.engine_list = [] # List of AsyncLLM instances for each stage
- 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"""
-
- # 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(request_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
-```
-
-### 4. Stage config for setting of different model stages
-```python
-@dataclass
-class StageConfig:
- """Configuration for a processing stage"""
- stage_id: int
- engine_type: str # "AR" or "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
-```
-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 = StageConfig()
-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 = StageConfig()
-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
-
-# 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 AsyncOmni(AsyncLLM):
- """Extended AsyncLLM supporting multiple engines and stage-based processing"""
-
- def __init__(self, stage_configs: List[StageConfig]):
- 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 Scheduler
-class DiffusionScheduler(Scheduler):
-
- 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. Diffusers Pipeline Executor (no worker)
-
-A single-process executor that directly runs the Diffusers pipeline without spawning workers or using RPC. Interfaces remain identical to `Executor` so the EngineCore loop is unchanged.
-
-#### Function map（Pipeline Executor）
-
-##### 1) Inherited and overridden
-```python
-from concurrent.futures import Future
-from typing import Optional, Union, Callable, Any
-import torch
-import torch.nn as nn
-from vllm.v1.executor.abstract import Executor
-from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
-from vllm.v1.outputs import ModelRunnerOutput
-from vllm.tasks import SupportedTask
-from diffusers import DiffusionPipeline
-
-class DiffusersPipelineExecutor(Executor):
- supports_pp: bool = False # Single-process, no TP/PP/DP
-
- def _init_executor(self) -> None:
- # Called by ExecutorBase.__init__
- self._failure_callback: Optional[Callable[[], None]] = None
- self._device = self._resolve_device()
- self._dtype = self._resolve_dtype()
- self._pipeline = self._build_pipeline(device=self._device, dtype=self._dtype)
- self._profiler = None
- self._is_failed = False
- self.is_sleeping = False
- self.sleeping_tags: set[str] = set()
-
- # major functions to build/run diffusers pipeline
- def _build_pipeline(self, device:torch.device, dtype:torch.dtype)->DiffusionPipeline:
- model_name = "Qwen/Qwen-Image"
-
- self.pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=dtype)
- self.pipe = pipe.to(device)
-
- def _run_pipeline(self, scheduler_output) -> ModelRunnerOutput: ...
- positive_magic = {
- "en": ", Ultra HD, 4K, cinematic composition.", # for english prompt
- "zh": ", 超清，4K，电影级构图." # for chinese prompt
- }
-
- # Generate image
- prompt_embeds = self._get_and_process_prompt_embeds(scheduler_output, positive_magic)
- negative_prompt_embeds = self.pipe.embed_prompt(" ")
-
- width, height = aspect_ratios["16:9"]
-
- image = pipe(
- prompt_embeds=prompt_embeds,
- negative_prompt_embeds=negative_prompt_embeds,
- width=width,
- height=height,
- num_inference_steps=50,
- true_cfg_scale=4.0,
- generator=torch.Generator(device="cuda").manual_seed(42)
- ).images[0]
-
- output = self.wrap_image_as_ModelRunnerOutput(image)
- return output
-
- # ---- Internal helpers (implementation-specific, not public API) ----
- def _resolve_device(self): ...
- def _resolve_dtype(self): ...
- def _get_model(self) -> nn.Module: ...
- def _get_and_process_prompt_embeds(self, scheduler_output, positive_magic):
- ...
- #append the positive_magic to prompt and embed them to prompt embed tensors
-
- # Functions related to workers should either raise NotImplementedError
- # or return empty defaults in this no-worker executor.
-
- def collective_rpc(self, method, timeout=None, args=(), kwargs=None) -> list[Any]:
- # No workers in pipeline executor
- raise NotImplementedError("No workers in DiffusersPipelineExecutor")
-
- def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
- return # no-op (pipeline already built in _init_executor)
-
- def register_failure_callback(self, callback):
- self._failure_callback = callback
-
- def determine_available_memory(self) -> list[int]: # bytes
- # Single device; return [available_bytes]. If CPU-only, return [0].
- return [self._determine_available_bytes(self._device)]
-
-```
-
-### 10. 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 `DiffusionModelRunner`; 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 DiffusionModelRunner.
-from vllm.v1.worker.gpu_worker import Worker as GPUWorker
-from vllm.v1.worker.diffusion_model_runner import DiffusionModelRunner
-
-class Worker(GPUWorker):
- def init_device(self) -> None:
- #those related to device check and init
- ...
-
- self.model_runner = DiffusionModelRunner(self.vllm_config, ...)
-```
-
-### 11. 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`：
-- DiffusionModelRunner: Use the `pooler_output=[Tensor,...]` to return multi modal tensors
-- ARModelRunner: 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 DiffusionModelRunner(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 ARModelRunner(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,...]
-```
-
-### 12. 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() → create request state
-
-AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() → create a background loop that continuously pulls output from EngineCore
-↓
-OutputProcessor.process_outputs() → update state and process outputs
-↓
-RequestState.make_request_output() → convert to the final output, formatted as RequestOutput or PoolingRequestOutput
-↓
-RequestOutputCollector.put() → push to the queue (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)
-```
-
-
-### 13. OutputProcessor
-
-For hidden state output, original OutputProcessor already support.
-
-We only need to add new one for final multimodal output.
-```python
-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)
-```
diff --git a/docs/source/architecture/ar-dit-main-architecture.png b/docs/source/architecture/ar-dit-main-architecture.png
new file mode 100644
index 000000000..200fddf2b
Binary files /dev/null and b/docs/source/architecture/ar-dit-main-architecture.png differ
diff --git a/docs/source/architecture/ar-main-architecture.png b/docs/source/architecture/ar-main-architecture.png
new file mode 100644
index 000000000..642c3c008
Binary files /dev/null and b/docs/source/architecture/ar-main-architecture.png differ
diff --git a/docs/source/architecture/dit-main-architecture.png b/docs/source/architecture/dit-main-architecture.png
new file mode 100644
index 000000000..97672e145
Binary files /dev/null and b/docs/source/architecture/dit-main-architecture.png differ
diff --git a/docs/source/architecture/vllm-omni-main-architecture.png b/docs/source/architecture/vllm-omni-main-architecture.png
new file mode 100644
index 000000000..98b7a9792
Binary files /dev/null and b/docs/source/architecture/vllm-omni-main-architecture.png differ
diff --git a/docs/source/architecture/vllm-omni-user-interface.png b/docs/source/architecture/vllm-omni-user-interface.png
new file mode 100644
index 000000000..867b3ae2f
Binary files /dev/null and b/docs/source/architecture/vllm-omni-user-interface.png differ