[Docs] optimize and supplement docs system (#86)

Signed-off-by: Gaohan123 <gaohan19@huawei.com>

vllm_omni/engine/__init__.py

@@ -49,6 +49,19 @@ class AdditionalInformationPayload(msgspec.Struct):
 
 
 class OmniEngineCoreRequest(EngineCoreRequest):
+    """Engine core request for omni models with embeddings support.
+
+    Extends the base EngineCoreRequest with support for prompt embeddings
+    and additional information payloads, enabling direct transfer of
+    pre-computed embeddings between pipeline stages.
+
+    Attributes:
+        prompt_embeds: Optional serialized prompt embeddings payload for
+            direct transfer between stages
+        additional_information: Optional serialized additional information
+            dictionary containing tensors or lists to pass along with the request
+    """
+
     # Optional prompt embeddings (direct-transfer version)
     prompt_embeds: Optional[PromptEmbedsPayload] = None
     # Optional additional information dictionary (serialized)

vllm_omni/engine/arg_utils.py

@@ -9,12 +9,33 @@ from vllm_omni.config import OmniModelConfig
 
 @dataclass
 class OmniEngineArgs(EngineArgs):
+    """Engine arguments for omni models, extending base EngineArgs.
+
+    Adds omni-specific configuration fields for multi-stage pipeline
+    processing and output type specification.
+
+    Args:
+        stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
+        model_stage: Stage type identifier, e.g., "thinker" or "talker"
+            (default: "thinker")
+        model_arch: Model architecture name
+            (default: "Qwen2_5OmniForConditionalGeneration")
+        engine_output_type: Optional output type specification for the engine.
+            Used to route outputs to appropriate processors (e.g., "image",
+            "audio", "latents"). If None, output type is inferred.
+    """
+
     stage_id: int = 0
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: Optional[str] = None
 
     def create_model_config(self) -> OmniModelConfig:
+        """Create an OmniModelConfig from these engine arguments.
+
+        Returns:
+            OmniModelConfig instance with all configuration fields set
+        """
         # First, get the base ModelConfig from the parent class
         base_config = super().create_model_config()
 
@@ -37,12 +58,33 @@ class OmniEngineArgs(EngineArgs):
 
 @dataclass
 class AsyncOmniEngineArgs(AsyncEngineArgs):
+    """Async engine arguments for omni models, extending base AsyncEngineArgs.
+
+    Adds omni-specific configuration fields for multi-stage pipeline
+    processing and output type specification in async contexts.
+
+    Args:
+        stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
+        model_stage: Stage type identifier, e.g., "thinker" or "talker"
+            (default: "thinker")
+        model_arch: Model architecture name
+            (default: "Qwen2_5OmniForConditionalGeneration")
+        engine_output_type: Optional output type specification for the engine.
+            Used to route outputs to appropriate processors (e.g., "image",
+            "audio", "latents"). If None, output type is inferred.
+    """
+
     stage_id: int = 0
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: Optional[str] = None
 
     def create_model_config(self) -> OmniModelConfig:
+        """Create an OmniModelConfig from these engine arguments.
+
+        Returns:
+            OmniModelConfig instance with all configuration fields set
+        """
         # First, get the base ModelConfig from the parent class
         base_config = super().create_model_config()
 

vllm_omni/engine/output_processor.py

@@ -19,6 +19,13 @@ logger = init_logger(__name__)
 
 
 class OmniRequestState(RequestState):
+    """Request state for omni models, tracking multimodal outputs.
+
+    Extends the base RequestState with support for accumulating
+    multimodal tensor outputs (e.g., images, audio, latents) that
+    are produced incrementally during generation.
+    """
+
     def __init__(
         self,
         *args,
@@ -86,6 +93,17 @@ class OmniRequestState(RequestState):
         )
 
     def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], mm_type: Optional[str]) -> None:
+        """Add a multimodal tensor to the accumulated outputs.
+
+        Accumulates multimodal tensors (e.g., image latents, audio waveforms)
+        that are produced incrementally during generation. Tensors are
+        concatenated along the first dimension.
+
+        Args:
+            tensor: Optional tensor to add. If None, this is a no-op.
+            mm_type: Optional modality type identifier (e.g., "image", "audio").
+                Used to categorize the output type.
+        """
         if tensor is None:
             return
         try:
@@ -115,6 +133,23 @@ class OmniRequestState(RequestState):
         stop_reason: Optional[Union[int, str]],
         kv_transfer_params: Optional[dict[str, Any]] = None,
     ) -> Optional[Union[OmniRequestOutput, PoolingRequestOutput]]:
+        """Create a request output from generation results.
+
+        Creates a RequestOutput or PoolingRequestOutput from the generated
+        tokens and accumulated multimodal outputs. Attaches multimodal
+        tensors to the completion output if available.
+
+        Args:
+            new_token_ids: List of newly generated token IDs
+            pooling_output: Optional pooling output tensor
+            finish_reason: Optional finish reason indicating why generation stopped
+            stop_reason: Optional stop reason (token ID or stop string)
+            kv_transfer_params: Optional KV cache transfer parameters
+
+        Returns:
+            OmniRequestOutput or PoolingRequestOutput if output should be
+            emitted (based on finish status and output kind), None otherwise
+        """
         finished = finish_reason is not None
         final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
 
@@ -179,6 +214,15 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         log_stats: bool,
         engine_core_output_type: Optional[str] = None,
     ):
+        """Initialize the multimodal output processor.
+
+        Args:
+            tokenizer: Tokenizer for detokenizing text outputs
+            log_stats: Whether to log statistics
+            engine_core_output_type: Optional output type specification
+                (e.g., "image", "audio", "latents"). Used to route outputs
+                to appropriate processors. If None, output type is inferred.
+        """
         super().__init__(tokenizer=tokenizer, log_stats=log_stats)
         self.output_handlers: dict[str, Callable[[EngineCoreOutput], None]] = {}
         self._reqid_to_mm_type: dict[str, str] = {}
@@ -186,6 +230,16 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         self.engine_core_output_type = engine_core_output_type
 
     def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None:
+        """Register a custom handler for a specific modality.
+
+        Allows custom processing logic for specific output modalities.
+        The handler is called before default processing for outputs
+        matching the specified modality.
+
+        Args:
+            modality: Modality name (e.g., "image", "audio", "latents")
+            handler: Callable that takes an EngineCoreOutput and processes it
+        """
         self.output_handlers[modality.lower()] = handler
 
     def add_request(
@@ -196,6 +250,21 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         request_index: int = 0,
         queue: Optional[RequestOutputCollector] = None,
     ) -> None:
+        """Add a new request to be processed.
+
+        Creates an OmniRequestState for the request and registers it
+        for output processing.
+
+        Args:
+            request: Engine core request to add
+            prompt: Optional prompt string for the request
+            parent_req: Optional parent request for parallel sampling
+            request_index: Index of the request in the batch
+            queue: Optional queue for collecting outputs
+
+        Raises:
+            ValueError: If the request ID is already registered
+        """
         request_id = request.request_id
         if request_id in self.request_states:
             raise ValueError(f"Request id {request_id} already running.")
@@ -220,6 +289,21 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         engine_core_timestamp: Optional[float] = None,
         iteration_stats: Optional[IterationStats] = None,
     ) -> OutputProcessorOutput:
+        """Process engine core outputs into request outputs.
+
+        Converts EngineCoreOutput objects into RequestOutput objects,
+        handling multimodal outputs by routing them through appropriate
+        processors and accumulating tensors in request states.
+
+        Args:
+            engine_core_outputs: List of engine core outputs to process
+            engine_core_timestamp: Optional timestamp for the outputs
+            iteration_stats: Optional iteration statistics
+
+        Returns:
+            OutputProcessorOutput containing processed request outputs
+            and list of request IDs to abort
+        """
         self._reqid_to_mm_type.clear()
         for eco in engine_core_outputs:
             mm_type = (self.engine_core_output_type or "").lower()

vllm_omni/engine/processor.py

@@ -30,8 +30,28 @@ logger = init_logger(__name__)
 
 
 class OmniProcessor(Processor):
+    """Processor for omni models, handling multimodal inputs and embeddings.
+
+    Extends the base vLLM Processor with support for processing prompt
+    embeddings and additional information payloads, enabling direct transfer
+    of pre-computed embeddings between pipeline stages.
+
+    Args:
+        vllm_config: Global vLLM configuration
+        tokenizer: Tokenizer instance for text processing
+        mm_registry: Multi-modal registry for processing multimodal inputs
+    """
+
     @staticmethod
     def _dtype_to_name(dtype: torch.dtype) -> str:
+        """Convert torch dtype to string representation.
+
+        Args:
+            dtype: PyTorch dtype to convert
+
+        Returns:
+            String representation of the dtype (e.g., "float32", "int64")
+        """
         mapping = {
             torch.float32: "float32",
             torch.float: "float32",
@@ -78,6 +98,35 @@ class OmniProcessor(Processor):
         priority: int = 0,
         data_parallel_rank: Optional[int] = None,
     ) -> tuple[Optional[str], OmniEngineCoreRequest]:
+        """Process input prompt into an engine core request.
+
+        Converts a prompt (text, tokens, or multimodal) into an
+        OmniEngineCoreRequest that can be processed by the engine.
+        Handles prompt embeddings and additional information payloads
+        for direct transfer between stages.
+
+        Args:
+            request_id: Unique identifier for this request
+            prompt: Input prompt (text, token IDs, embeddings, or multimodal)
+            params: Sampling or pooling parameters for generation
+            arrival_time: Optional arrival timestamp (defaults to current time)
+            lora_request: Optional LoRA adapter request
+            tokenization_kwargs: Optional additional tokenization arguments
+            trace_headers: Optional tracing headers for observability
+            priority: Request priority (higher values processed first)
+            data_parallel_rank: Optional data parallel rank for distributed
+                inference
+
+        Returns:
+            Tuple of (prompt_string, OmniEngineCoreRequest) where:
+                - prompt_string: The original prompt as a string, or None if
+                  using embeddings
+                - OmniEngineCoreRequest: Processed request ready for the engine
+
+        Raises:
+            ValueError: If data_parallel_rank is out of range or prompt_embeds
+                has incorrect shape
+        """
         # TODO(woosuk): Support pooling models.
         # TODO(woosuk): Support encoder-decoder models.
         self._validate_lora(lora_request)