chore: flux benchmarking script + code clean (ai-dynamo#8083)

Signed-off-by: ayushag <ayushag@nvidia.com>

Author: ayushag-nv

benchmarks/omni/image/aiperf_image_gen.sh

@@ -0,0 +1,79 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Generic aiperf benchmark for vLLM-Omni text-to-image generation.
+# Assumes the server (Dynamo or native vllm-omni) is already running.
+#
+# Usage:
+#   bash aiperf_image_gen.sh [OPTIONS]
+#
+# Options:
+#   --model <model>           Model to benchmark (default: black-forest-labs/FLUX.2-klein-4B)
+#   --url <url>               Server URL (default: http://localhost:8000)
+#   --concurrency <n>         Number of concurrent requests (default: 1)
+#   --request-count <n>       Total requests to send (default: 10)
+#   --warmup-count <n>        Warmup requests before measurement (default: 2)
+#   --image-size <WxH>        Generated image size (default: 1024x1024)
+#   --response-format <fmt>   Response format: url or b64_json (default: url)
+#   --prompt-tokens-mean <n>  Mean synthetic prompt length in tokens (default: 50)
+#   --prompt-tokens-stddev <n> Stddev of synthetic prompt length (default: 10)
+#   -h, --help                Show this help message
+#
+# Examples:
+#   bash aiperf_image_gen.sh
+#   bash aiperf_image_gen.sh --model zai-org/GLM-Image --concurrency 4
+#   bash aiperf_image_gen.sh --model Qwen/Qwen-Image --image-size 512x512 --request-count 20
+
+MODEL="black-forest-labs/FLUX.2-klein-4B"
+URL="http://localhost:8000"
+CONCURRENCY=1
+REQUEST_COUNT=10
+WARMUP_COUNT=2
+IMAGE_SIZE="1024x1024"
+RESPONSE_FORMAT="url"
+PROMPT_TOKENS_MEAN=50
+PROMPT_TOKENS_STDDEV=10
+ARTIFACT_DIR=""
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --model)                MODEL=$2;               shift 2 ;;
+        --url)                  URL=$2;                 shift 2 ;;
+        --concurrency)          CONCURRENCY=$2;         shift 2 ;;
+        --request-count)        REQUEST_COUNT=$2;       shift 2 ;;
+        --warmup-count)         WARMUP_COUNT=$2;        shift 2 ;;
+        --image-size)           IMAGE_SIZE=$2;          shift 2 ;;
+        --response-format)      RESPONSE_FORMAT=$2;     shift 2 ;;
+        --prompt-tokens-mean)   PROMPT_TOKENS_MEAN=$2;  shift 2 ;;
+        --prompt-tokens-stddev) PROMPT_TOKENS_STDDEV=$2; shift 2 ;;
+        --artifact-dir)         ARTIFACT_DIR=$2;        shift 2 ;;
+        -h|--help)
+            sed -n '/^# Usage/,/^[^#]/p' "$0" | grep '^#' | sed 's/^# \?//'
+            exit 0 ;;
+        *) echo "Unknown option: $1"; exit 1 ;;
+    esac
+done
+
+AIPERF_ARGS=(
+    aiperf profile
+    --model "$MODEL"
+    --tokenizer gpt2
+    --url "$URL"
+    --endpoint-type image-generation
+    --synthetic-input-tokens-mean "$PROMPT_TOKENS_MEAN"
+    --synthetic-input-tokens-stddev "$PROMPT_TOKENS_STDDEV"
+    --extra-inputs "size:${IMAGE_SIZE}"
+    --extra-inputs "response_format:${RESPONSE_FORMAT}"
+    --concurrency "$CONCURRENCY"
+    --request-count "$REQUEST_COUNT"
+    --warmup-request-count "$WARMUP_COUNT"
+    --ui none
+    --no-server-metrics
+)
+
+if [[ -n "$ARTIFACT_DIR" ]]; then
+    AIPERF_ARGS+=(--artifact-dir "$ARTIFACT_DIR")
+fi
+
+"${AIPERF_ARGS[@]}"

components/src/dynamo/vllm/omni/args.py

@@ -4,6 +4,7 @@
 """Omni-specific argument parsing for python -m dynamo.vllm.omni."""
 
 import argparse
+import dataclasses
 import logging
 from typing import Optional
 
@@ -24,15 +25,50 @@
 logger = logging.getLogger(__name__)
 
 
-class OmniArgGroup(ArgGroup):
-    """Diffusion pipeline kwargs passed through to AsyncOmni() constructor.
+@dataclasses.dataclass
+class OmniDiffusionKwargs:
+    """AsyncOmni constructor kwargs for diffusion engine configuration.
+
+    Every field here is passed directly to AsyncOmni(**kwargs) and consumed by
+    _create_default_diffusion_stage_cfg() in vllm-omni. Adding a new vllm-omni
+    diffusion flag only requires adding it here and to OmniArgGroup — the
+    passthrough in base_handler is automatic.
+    """
+
+    enable_layerwise_offload: bool = False
+    layerwise_num_gpu_layers: int = 1
+    vae_use_slicing: bool = False
+    vae_use_tiling: bool = False
+    boundary_ratio: float = 0.875
+    flow_shift: Optional[float] = None
+    cache_backend: Optional[str] = None
+    cache_config: Optional[str] = None
+    enable_cache_dit_summary: bool = False
+    enable_cpu_offload: bool = False
+    enforce_eager: bool = False
+
 
-    These are NOT part of OmniEngineArgs (which handles vLLM engine-level
-    args like model, tp, max_model_len). Instead they are direct constructor
-    kwargs for AsyncOmni and need Dynamo-side env-var (DYN_OMNI_*) support,
-    so we define them here rather than relying on the upstream arg parser.
+@dataclasses.dataclass
+class OmniParallelKwargs:
+    """Diffusion parallelism configuration passed to DiffusionParallelConfig.
+
+    Every field here maps 1:1 to a DiffusionParallelConfig field (excluding
+    tensor_parallel_size which comes from engine_args, and fixed/derived fields).
+    Adding a new parallelism field only requires adding it here and to OmniArgGroup.
     """
 
+    ulysses_degree: int = 1
+    ring_degree: int = 1
+    cfg_parallel_size: int = 1
+    vae_patch_parallel_size: int = 1
+    use_hsdp: bool = False
+    hsdp_shard_size: int = -1
+    hsdp_replicate_size: int = 1
+
+
+class OmniArgGroup(ArgGroup):
+    """CLI argument definitions for Dynamo vLLM-Omni."""
+
     name = "dynamo-omni"
 
     def add_arguments(self, parser) -> None:
@@ -49,7 +85,6 @@ def add_arguments(self, parser) -> None:
             help="Path to vLLM-Omni stage configuration YAML file (optional).",
         )
 
-        # Video encoding
         add_argument(
             g,
             flag_name="--default-video-fps",
@@ -59,7 +94,7 @@ def add_arguments(self, parser) -> None:
             help="Default frames per second for generated videos.",
         )
 
-        # Layerwise offloading
+        # OmniDiffusionKwargs fields
         add_negatable_bool_argument(
             g,
             flag_name="--enable-layerwise-offload",
@@ -75,8 +110,6 @@ def add_arguments(self, parser) -> None:
             arg_type=int,
             help="Number of ready layers (blocks) to keep on GPU during generation.",
         )
-
-        # VAE optimization
         add_negatable_bool_argument(
             g,
             flag_name="--vae-use-slicing",
@@ -91,8 +124,6 @@ def add_arguments(self, parser) -> None:
             default=False,
             help="Enable VAE tiling for memory optimization in diffusion models.",
         )
-
-        # Diffusion scheduling
         add_argument(
             g,
             flag_name="--boundary-ratio",
@@ -113,8 +144,6 @@ def add_arguments(self, parser) -> None:
             arg_type=float,
             help="Scheduler flow_shift parameter (5.0 for 720p, 12.0 for 480p).",
         )
-
-        # Cache acceleration
         add_argument(
             g,
             flag_name="--cache-backend",
@@ -141,8 +170,6 @@ def add_arguments(self, parser) -> None:
             default=False,
             help="Enable cache-dit summary logging after diffusion forward passes.",
         )
-
-        # Execution mode
         add_negatable_bool_argument(
             g,
             flag_name="--enable-cpu-offload",
@@ -204,7 +231,7 @@ def add_arguments(self, parser) -> None:
             help="Maximum size in bytes for reference audio files (default: 50MB).",
         )
 
-        # Diffusion parallel configuration
+        # OmniParallelKwargs fields
         add_argument(
             g,
             flag_name="--ulysses-degree",
@@ -227,9 +254,43 @@ def add_arguments(self, parser) -> None:
             env_var="DYN_OMNI_CFG_PARALLEL_SIZE",
             default=1,
             arg_type=int,
-            choices=[1, 2],
+            choices=[1, 2, 3],
             help="Number of GPUs used for classifier free guidance parallelism.",
         )
+        add_argument(
+            g,
+            flag_name="--vae-patch-parallel-size",
+            env_var="DYN_OMNI_VAE_PATCH_PARALLEL_SIZE",
+            default=1,
+            arg_type=int,
+            help="Number of ranks used for VAE patch/tile parallelism during decode/encode.",
+        )
+        add_negatable_bool_argument(
+            g,
+            flag_name="--use-hsdp",
+            env_var="DYN_OMNI_USE_HSDP",
+            default=False,
+            help=(
+                "Enable Hybrid Sharded Data Parallel (HSDP) for diffusion models. "
+                "Shards model weights across GPUs to reduce per-GPU memory usage."
+            ),
+        )
+        add_argument(
+            g,
+            flag_name="--hsdp-shard-size",
+            env_var="DYN_OMNI_HSDP_SHARD_SIZE",
+            default=-1,
+            arg_type=int,
+            help="Number of GPUs to shard model weights across when using HSDP (-1 = auto).",
+        )
+        add_argument(
+            g,
+            flag_name="--hsdp-replicate-size",
+            env_var="DYN_OMNI_HSDP_REPLICATE_SIZE",
+            default=1,
+            arg_type=int,
+            help="Number of HSDP replica groups (default: 1).",
+        )
 
         # Disaggregated stage worker flags
         add_argument(
@@ -244,7 +305,6 @@ def add_arguments(self, parser) -> None:
                 "Requires --stage-configs-path."
             ),
         )
-
         add_negatable_bool_argument(
             g,
             flag_name="--omni-router",
@@ -263,30 +323,18 @@ class OmniConfig(DynamoRuntimeConfig):
     component: str = "backend"
     endpoint: Optional[str] = None
 
-    # mirror vLLM
     model: str
     served_model_name: Optional[str] = None
-
-    # vLLM-Omni engine args
     engine_args: OmniEngineArgs
 
-    # OmniArgGroup fields (populated by from_cli_args)
     stage_configs_path: Optional[str] = None
     default_video_fps: int = 16
-    enable_layerwise_offload: bool = False
-    layerwise_num_gpu_layers: int = 1
-    vae_use_slicing: bool = False
-    vae_use_tiling: bool = False
-    boundary_ratio: float = 0.875
-    flow_shift: Optional[float] = None
-    cache_backend: Optional[str] = None
-    cache_config: Optional[str] = None
-    enable_cache_dit_summary: bool = False
-    enable_cpu_offload: bool = False
-    enforce_eager: bool = False
-    ulysses_degree: int = 1
-    ring_degree: int = 1
-    cfg_parallel_size: int = 1
+
+    # Nested structs — each group of fields has a clear destination
+    diffusion: OmniDiffusionKwargs = dataclasses.field(
+        default_factory=OmniDiffusionKwargs
+    )
+    parallel: OmniParallelKwargs = dataclasses.field(default_factory=OmniParallelKwargs)
 
     # TTS parameters
     tts_max_instructions_length: int = 500
@@ -299,15 +347,36 @@ class OmniConfig(DynamoRuntimeConfig):
     stage_id: Optional[int] = None
     omni_router: bool = False
 
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "OmniConfig":
+        config = super().from_cli_args(args)
+        config.diffusion = dataclasses.replace(
+            OmniDiffusionKwargs(),
+            **{
+                f.name: getattr(args, f.name)
+                for f in dataclasses.fields(OmniDiffusionKwargs)
+                if hasattr(args, f.name)
+            },
+        )
+        config.parallel = dataclasses.replace(
+            OmniParallelKwargs(),
+            **{
+                f.name: getattr(args, f.name)
+                for f in dataclasses.fields(OmniParallelKwargs)
+                if hasattr(args, f.name)
+            },
+        )
+        return config
+
     def validate(self) -> None:
         DynamoRuntimeConfig.validate(self)
         if self.default_video_fps <= 0:
             raise ValueError("--default-video-fps must be > 0")
-        if self.ulysses_degree <= 0:
+        if self.parallel.ulysses_degree <= 0:
             raise ValueError("--ulysses-degree must be > 0")
-        if self.ring_degree <= 0:
+        if self.parallel.ring_degree <= 0:
             raise ValueError("--ring-degree must be > 0")
-        if not (0 < self.boundary_ratio <= 1):
+        if not (0 < self.diffusion.boundary_ratio <= 1):
             raise ValueError("--boundary-ratio must be in (0, 1]")
         if self.stage_configs_path is None:
             if self.stage_id is not None:
@@ -334,7 +403,6 @@ def parse_omni_args() -> OmniConfig:
     dynamo_runtime_argspec.add_arguments(parser)
     omni_argspec.add_arguments(parser)
 
-    # Add vLLM-Omni engine args
     vg = parser.add_argument_group(
         "vLLM-Omni Engine Options. Please refer to vLLM-Omni documentation for more details."
     )
@@ -349,7 +417,6 @@ def parse_omni_args() -> OmniConfig:
     args, unknown = parser.parse_known_args()
     config = OmniConfig.from_cli_args(args)
 
-    # Default endpoint to "generate" if not explicitly provided by user
     if config.endpoint is None:
         config.endpoint = "generate"
 

components/src/dynamo/vllm/omni/base_handler.py

@@ -4,6 +4,7 @@
 """Base handler for vLLM-Omni multi-stage pipelines."""
 
 import asyncio
+import dataclasses
 import logging
 import time
 from typing import Any, AsyncGenerator, Dict
@@ -74,31 +75,17 @@ def _build_omni_kwargs(self, config) -> Dict[str, Any]:
         if config.stage_configs_path:
             omni_kwargs["stage_configs_path"] = config.stage_configs_path
 
-        # Diffusion engine-level params — read directly from config namespace
-        diffusion_fields = [
-            "enable_layerwise_offload",
-            "layerwise_num_gpu_layers",
-            "vae_use_slicing",
-            "vae_use_tiling",
-            "boundary_ratio",
-            "flow_shift",
-            "cache_backend",
-            "cache_config",
-            "enable_cache_dit_summary",
-            "enable_cpu_offload",
-            "enforce_eager",
-        ]
-        for field in diffusion_fields:
-            value = getattr(config, field, None)
+        for field, value in dataclasses.asdict(config.diffusion).items():
             if value is not None:
                 omni_kwargs[field] = value
 
-        # Build DiffusionParallelConfig if available
+        # tensor_parallel_size comes from engine_args (vLLM's --tensor-parallel-size)
         if DiffusionParallelConfig is not None:
             parallel_config = DiffusionParallelConfig(
-                ulysses_degree=getattr(config, "ulysses_degree", 1),
-                ring_degree=getattr(config, "ring_degree", 1),
-                cfg_parallel_size=getattr(config, "cfg_parallel_size", 1),
+                tensor_parallel_size=getattr(
+                    config.engine_args, "tensor_parallel_size", 1
+                ),
+                **dataclasses.asdict(config.parallel),
             )
             omni_kwargs["parallel_config"] = parallel_config
         else: