[diffusion] chore: clean excessive document (#16986)

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -633,8 +633,6 @@ def patch_tensor_parallel_group(tp_group: GroupCoordinator):
     This method is for draft workers of speculative decoding to run draft model
     with different tp degree from that of target model workers.
 
-    Args:
-        tp_group (GroupCoordinator): the tp group coordinator
     """
     global _TP_STATE_PATCHED
     assert not _TP_STATE_PATCHED, "Should not call when it's already patched"

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -87,12 +87,6 @@ class DiffGenerator:
         """
         Create a DiffGenerator from a pretrained model.
 
-        Args:
-            **kwargs: Additional arguments to customize model loading, set any ServerArgs or PipelineConfig attributes here.
-
-        Returns:
-            The created DiffGenerator
-
         Priority level: Default pipeline config < User's pipeline config < User's kwargs
         """
         # If users also provide some kwargs, it will override the ServerArgs and PipelineConfig.

python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py

@@ -139,12 +139,6 @@ class AttentionImpl(ABC, Generic[T]):
 
         Called AFTER all_to_all for distributed attention
 
-        Args:
-            qkv: The query-key-value tensor
-            attn_metadata: Metadata for the attention operation
-
-        Returns:
-            Processed QKV tensor
         """
         return qkv
 
@@ -161,12 +155,6 @@ class AttentionImpl(ABC, Generic[T]):
 
         Called BEFORE all_to_all for distributed attention
 
-        Args:
-            output: The output tensor from the attention operation
-            attn_metadata: Metadata for the attention operation
-
-        Returns:
-            Postprocessed output tensor
         """
 
         return output

python/sglang/multimodal_gen/runtime/layers/attention/selector.py

@@ -170,7 +170,6 @@ def global_force_attn_backend_context_manager(
     manager.
 
     Arguments:
-
     * attn_backend: attention backend to force
 
     Returns:

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -254,9 +254,6 @@ class ComponentLoader(ABC):
         Args:
             module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler")
             transformers_or_diffusers: Whether the module is from transformers or diffusers
-
-        Returns:
-            A component loader for the specified module type
         """
         # Map of module types to their loader classes and expected library
         module_type = _normalize_module_type(module_type)
@@ -803,8 +800,6 @@ class PipelineComponentLoader:
             component_model_path: Path to the component model
             transformers_or_diffusers: Whether the module is from transformers or diffusers
 
-        Returns:
-            The loaded module
         """
 
         # Get the appropriate loader for this module type

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -181,8 +181,6 @@ def shard_model(
             which modules to shard with FSDP.
         pin_cpu_memory (bool): If set to True, FSDP will pin the CPU memory of the offloaded parameters.
 
-    Raises:
-        ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered.
     """
     if fsdp_shard_conditions is None or len(fsdp_shard_conditions) == 0:
         logger.warning(
@@ -244,8 +242,6 @@ def load_model_from_full_model_state_dict(
             * **missing_keys** is a list of str containing the missing keys
             * **unexpected_keys** is a list of str containing the unexpected keys
 
-    Raises:
-        NotImplementedError: If got FSDP with more than 1D.
     """
     meta_sd = model.state_dict()
     param_dict = dict(model.named_parameters())

python/sglang/multimodal_gen/runtime/models/encoders/clip.py

@@ -447,10 +447,6 @@ class CLIPTextTransformer(nn.Module):
         inputs_embeds: torch.Tensor | None = None,
         output_hidden_states: bool | None = None,
     ) -> BaseEncoderOutput:
-        r"""
-        Returns:
-
-        """
         output_hidden_states = (
             output_hidden_states
             if output_hidden_states is not None

python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -196,16 +196,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler
     ) -> torch.FloatTensor:
         """
         Forward process in flow-matching
-
-        Args:
-            sample (`torch.FloatTensor`):
-                The input sample.
-            timestep (`int`, *optional*):
-                The current timestep in the diffusion chain.
-
-        Returns:
-            `torch.FloatTensor`:
-                A scaled input sample.
         """
         # Make sure sigmas and timesteps have the same device and dtype as original_samples
         sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)

python/sglang/multimodal_gen/runtime/models/utils.py

@@ -78,16 +78,7 @@ def modulate(
     shift: torch.Tensor | None = None,
     scale: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """modulate by shift and scale
-
-    Args:
-        x (torch.Tensor): input tensor.
-        shift (torch.Tensor, optional): shift tensor. Defaults to None.
-        scale (torch.Tensor, optional): scale tensor. Defaults to None.
-
-    Returns:
-        torch.Tensor: the output tensor after modulate.
-    """
+    """modulate by shift and scale"""
     if scale is None and shift is None:
         return x
     elif shift is None:

python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py

@@ -544,8 +544,6 @@ class AutoencoderKL(nn.Module):
             sample (`torch.Tensor`): Input sample.
             sample_posterior (`bool`, *optional*, defaults to `False`):
                 Whether to sample from the posterior.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
         """
         x = sample
         posterior = self.encode(x).latent_dist

python/sglang/multimodal_gen/runtime/models/vision_utils.py

@@ -101,10 +101,6 @@ def load_image(
         convert_method (Callable[[PIL.Image.Image], PIL.Image.Image], *optional*):
             A conversion method to apply to the image after loading it. When set to `None` the image will be converted
             "RGB".
-
-    Returns:
-        `PIL.Image.Image`:
-            A PIL Image.
     """
     if isinstance(image, str):
         if image.startswith("http://") or image.startswith("https://"):

python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py

@@ -82,13 +82,6 @@ class PipelineStage(ABC):
                 result.add_check("image_latent", batch.image_latent, V.is_tensor)
                 return result
 
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
-
-        Returns:
-            A VerificationResult containing the verification status.
-
         """
         # Default implementation - no verification
         return VerificationResult()
@@ -119,9 +112,7 @@ class PipelineStage(ABC):
         """
         Verify the output for the stage.
 
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
+
 
         Returns:
             A VerificationResult containing the verification status.
@@ -182,9 +173,7 @@ class PipelineStage(ABC):
         Execute the stage's processing on the batch with optional verification and logging.
         Should not be overridden by subclasses.
 
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
+
 
         Returns:
             The updated batch information after this stage's processing.
@@ -232,9 +221,7 @@ class PipelineStage(ABC):
         This method should be implemented by subclasses to provide the forward
         processing logic for the stage.
 
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
+
 
         Returns:
             The updated batch information after this stage's processing.

python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py

@@ -38,9 +38,7 @@ class ConditioningStage(PipelineStage):
         """
         Apply conditioning to the diffusion process.
 
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
+
 
         Returns:
             The batch with applied conditioning.

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py

@@ -193,21 +193,6 @@ class DecodingStage(PipelineStage):
         representations to pixel-space video/images. It also optionally decodes
         trajectory latents for visualization purposes.
 
-        Args:
-            batch: The current batch containing:
-                - latents: Tensor to decode (batch, channels, frames, height_latents, width_latents)
-                - return_trajectory_decoded (optional): Flag to decode trajectory latents
-                - trajectory_latents (optional): Latents at different timesteps
-                - trajectory_timesteps (optional): Corresponding timesteps
-            server_args: Configuration containing:
-                - vae_cpu_offload: Whether to offload VAE to CPU after decoding
-                - model_loaded: Track VAE loading state
-                - model_paths: Path to VAE model if loading needed
-
-        Returns:
-            Modified batch with:
-                - output: Decoded frames (batch, channels, frames, height, width) as CPU float32
-                - trajectory_decoded (if requested): List of decoded frames per timestep
         """
         # load vae if not already loaded (used for memory constrained devices)
         self.load_model()

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -478,10 +478,6 @@ class DenoisingStage(PipelineStage):
         """
         Prepare all necessary invariant variables for the denoising loop.
 
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
-
         Returns:
             A dictionary containing all the prepared variables for the denoising loop.
         """
@@ -940,13 +936,6 @@ class DenoisingStage(PipelineStage):
     ) -> Req:
         """
         Run the denoising loop.
-
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
-
-        Returns:
-            The batch with denoised latents.
         """
         # Prepare variables for the denoising loop
 
@@ -1127,13 +1116,6 @@ class DenoisingStage(PipelineStage):
     ) -> tqdm:
         """
         Create a progress bar for the denoising process.
-
-        Args:
-            iterable: The iterable to iterate over.
-            total: The total number of items.
-
-        Returns:
-            A tqdm progress bar.
         """
         local_rank = get_world_group().local_rank
         disable = local_rank != 0
@@ -1176,11 +1158,6 @@ class DenoisingStage(PipelineStage):
 
         Args:
             i: The current timestep index.
-            batch: The current batch information.
-            server_args: The inference arguments.
-
-        Returns:
-            The attention metadata, or None if not applicable.
         """
         attn_metadata = None
         self.attn_metadata_builder = None
@@ -1375,10 +1352,6 @@ class DenoisingStage(PipelineStage):
     def prepare_sta_param(self, batch: Req, server_args: ServerArgs):
         """
         Prepare Sliding Tile Attention (STA) parameters and settings.
-
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
         """
         # TODO(kevin): STA mask search, currently only support Wan2.1 with 69x768x1280
         STA_mode = server_args.STA_mode

python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py

@@ -60,9 +60,7 @@ class EncodingStage(PipelineStage):
         """
         Encode pixel space representations into latent space.
 
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
+
 
         Returns:
             The batch with encoded latents.

python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py

@@ -96,13 +96,6 @@ class ImageEncodingStage(PipelineStage):
     ) -> Req:
         """
         Encode the prompt into image encoder hidden states.
-
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
-
-        Returns:
-            The batch with encoded prompt embeddings.
         """
 
         if batch.condition_image is None:
@@ -212,13 +205,6 @@ class ImageVAEEncodingStage(PipelineStage):
     ) -> Req:
         """
         Encode pixel representations into latent space.
-
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
-
-        Returns:
-            The batch with encoded outputs.
         """
 
         if batch.condition_image is None:

python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py

@@ -182,13 +182,6 @@ class InputValidationStage(PipelineStage):
     ) -> Req:
         """
         Validate and prepare inputs.
-
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
-
-        Returns:
-            The validated batch information.
         """
 
         self._generate_seeds(batch, server_args)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py

@@ -42,9 +42,7 @@ class LatentPreparationStage(PipelineStage):
         """
         Prepare initial latent variables for the diffusion process.
 
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
+
 
         Returns:
             The batch with prepared latent variables.
@@ -111,9 +109,7 @@ class LatentPreparationStage(PipelineStage):
         """
         Adjust video length based on VAE version.
 
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
+
 
         Returns:
             The batch with adjusted video length.

python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py

@@ -53,13 +53,6 @@ class TextEncodingStage(PipelineStage):
     ) -> Req:
         """
         Encode the prompt into text encoder hidden states.
-
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
-
-        Returns:
-            The batch with encoded prompt embeddings.
         """
         assert len(self.tokenizers) == len(self.text_encoders)
         assert len(self.text_encoders) == len(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py

@@ -61,9 +61,7 @@ class TimestepPreparationStage(PipelineStage):
         """
         Prepare timesteps for the diffusion process.
 
-        Args:
-            batch: The current batch information.
-            server_args: The inference arguments.
+
 
         Returns:
             The batch with prepared timesteps.

python/sglang/multimodal_gen/runtime/server_args.py

@@ -738,17 +738,6 @@ class ServerArgs:
     ) -> int:
         """
         Find an available port with retry logic.
-
-        Args:
-            port: Initial port to check
-            port_inc: Port increment for each attempt
-            max_attempts: Maximum number of attempts to find an available port
-
-        Returns:
-            An available port number
-
-        Raises:
-            RuntimeError: If no available port is found after max_attempts
         """
         attempts = 0
         original_port = port
@@ -1088,13 +1077,6 @@ _global_server_args = None
 def prepare_server_args(argv: list[str]) -> ServerArgs:
     """
     Prepare the inference arguments from the command line arguments.
-
-    Args:
-        argv: The command line arguments. Typically, it should be `sys.argv[1:]`
-            to ensure compatibility with `parse_args` when no arguments are passed.
-
-    Returns:
-        The inference arguments.
     """
     parser = FlexibleArgumentParser()
     ServerArgs.add_cli_args(parser)

python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py

@@ -296,14 +296,7 @@ def load_dict(file_path):
 def get_diffusers_component_config(
     model_path: str,
 ) -> dict[str, Any]:
-    """Gets a configuration of a submodule for the given diffusers model.
-
-    Args:
-        model_path: the path of the submodule (can be local path or HuggingFace model ID)
-
-    Returns:
-        The loaded configuration.
-    """
+    """Gets a configuration of a submodule for the given diffusers model."""
 
     # Download from HuggingFace Hub if path doesn't exist locally
     if not os.path.exists(model_path):