feat(vllm-omni): pass model-specific video parameters

components/src/dynamo/common/tests/test_video_utils.py

@@ -20,6 +20,20 @@ def make_frames(n=3, h=8, w=8) -> np.ndarray:
     return np.zeros((n, h, w, 3), dtype=np.uint8)
 
 
+# ---------------------------------------------------------------------------
+# compute_num_frames
+# ---------------------------------------------------------------------------
+
+
+def test_compute_num_frames_rounds_fractional_frame_rate():
+    from dynamo.common.utils.video_utils import compute_num_frames
+
+    result = compute_num_frames(seconds=10, fps=23.976)
+
+    assert result == 240
+    assert isinstance(result, int)
+
+
 # ---------------------------------------------------------------------------
 # encode_to_video_bytes
 # ---------------------------------------------------------------------------

components/src/dynamo/common/utils/video_utils.py

@@ -45,7 +45,7 @@ def parse_size(
 def compute_num_frames(
     num_frames: int | None = None,
     seconds: int | None = None,
-    fps: int | None = None,
+    fps: float | None = None,
     default_fps: int = DEFAULT_VIDEO_FPS,
     default_num_frames: int = DEFAULT_VIDEO_NUM_FRAMES,
 ) -> int:
@@ -58,7 +58,7 @@ def compute_num_frames(
     if seconds is not None or fps is not None:
         _seconds = seconds if seconds is not None else 4
         _fps = fps if fps is not None else default_fps
-        return _seconds * _fps
+        return round(_seconds * _fps)
     return default_num_frames
 
 

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

@@ -56,6 +56,9 @@ class OmniDiffusionKwargs:
     cache_config: Optional[str] = None
     enable_cache_dit_summary: bool = False
     enable_cpu_offload: bool = False
+    task_type: Optional[str] = None
+    lora_path: Optional[list[str]] = None
+    diffusion_attention_backend: Optional[str] = None
     enforce_eager: bool = False
 
 
@@ -69,6 +72,7 @@ class OmniParallelKwargs:
     """
 
     ulysses_degree: int = 1
+    ulysses_a2a_permute: bool = False
     ring_degree: int = 1
     allgather_degree: int = 1
     cfg_parallel_size: int = 1
@@ -191,6 +195,32 @@ def add_arguments(self, parser) -> None:
             default=False,
             help="Enable CPU offloading for diffusion models to reduce GPU memory usage.",
         )
+        add_argument(
+            g,
+            flag_name="--task-type",
+            env_var="DYN_OMNI_TASK_TYPE",
+            default=None,
+            help="Model-defined task or checkpoint partition selected at startup.",
+        )
+        add_argument(
+            g,
+            flag_name="--lora-path",
+            env_var="DYN_OMNI_LORA_PATH",
+            default=None,
+            nargs="+",
+            env_value_type=list,
+            help=(
+                "Diffusion checkpoint adapter path(s) fused by vLLM-Omni at "
+                "startup. This is separate from request-time LoRA loading."
+            ),
+        )
+        add_argument(
+            g,
+            flag_name="--diffusion-attention-backend",
+            env_var="DYN_OMNI_DIFFUSION_ATTENTION_BACKEND",
+            default=None,
+            help="vLLM-Omni diffusion attention backend.",
+        )
         add_negatable_bool_argument(
             g,
             flag_name="--enforce-eager",
@@ -254,6 +284,17 @@ def add_arguments(self, parser) -> None:
             arg_type=int,
             help="Number of GPUs used for Ulysses sequence parallelism in diffusion.",
         )
+        add_negatable_bool_argument(
+            g,
+            flag_name="--ulysses-a2a-permute",
+            env_var="DYN_OMNI_ULYSSES_A2A_PERMUTE",
+            default=False,
+            help=(
+                "Use vLLM-Omni's fused permute-free all-to-all for eligible "
+                "Ulysses exchanges. Requires --ulysses-degree > 1; enabling "
+                "it JIT-compiles a CUDA kernel at worker start."
+            ),
+        )
         add_argument(
             g,
             flag_name="--ring-degree",

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

@@ -146,6 +146,41 @@ def _build_omni_kwargs(self, config) -> Dict[str, Any]:
         # source in engine_args while forwarding them to Omni's diffusion
         # topology explicitly.
         if DiffusionParallelConfig is not None:
+            # Compatibility workaround for the XPU image's vLLM-Omni 0.27 pin
+            # in container/context.yaml. Remove this filtering when that pin is
+            # bumped to a version supporting every OmniParallelKwargs field.
+            parallel_kwargs = dataclasses.asdict(config.parallel)
+            supported_parallel_fields = {
+                field.name for field in dataclasses.fields(DiffusionParallelConfig)
+            }
+            unsupported_parallel_fields = sorted(
+                set(parallel_kwargs) - supported_parallel_fields
+            )
+            if unsupported_parallel_fields:
+                default_parallel_kwargs = dataclasses.asdict(type(config.parallel)())
+                unsupported_non_defaults = [
+                    field
+                    for field in unsupported_parallel_fields
+                    if parallel_kwargs[field] != default_parallel_kwargs[field]
+                ]
+                if unsupported_non_defaults:
+                    options = ", ".join(unsupported_non_defaults)
+                    raise ValueError(
+                        "Installed vLLM-Omni does not support non-default "
+                        f"parallel option(s): {options}. Upgrade vLLM-Omni or "
+                        "remove the unsupported option(s)."
+                    )
+                logger.debug(
+                    "Ignoring parallel options unavailable in the installed "
+                    "vLLM-Omni: %s",
+                    ", ".join(unsupported_parallel_fields),
+                )
+                parallel_kwargs = {
+                    field: value
+                    for field, value in parallel_kwargs.items()
+                    if field in supported_parallel_fields
+                }
+
             parallel_config = DiffusionParallelConfig(
                 tensor_parallel_size=getattr(
                     config.engine_args, "tensor_parallel_size", 1
@@ -154,7 +189,7 @@ def _build_omni_kwargs(self, config) -> Dict[str, Any]:
                     config.engine_args, "pipeline_parallel_size", 1
                 ),
                 data_parallel_size=getattr(config.engine_args, "data_parallel_size", 1),
-                **dataclasses.asdict(config.parallel),
+                **parallel_kwargs,
             )
             omni_kwargs["parallel_config"] = parallel_config
         else:

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

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 import asyncio
+import copy
 import functools
 import logging
 import os
@@ -35,7 +36,11 @@
     get_output_modalities,
     parse_request_type,
 )
-from dynamo.common.utils.video_utils import compute_num_frames, parse_size
+from dynamo.common.utils.video_utils import (
+    DEFAULT_VIDEO_NUM_FRAMES,
+    compute_num_frames,
+    parse_size,
+)
 from dynamo.llm import (
     ModelInput,
     ModelRuntimeConfig,
@@ -683,6 +688,95 @@ def _build_sampling_params_list(
                 )
         return result if result else [diffusion_sp]
 
+    def _build_video_sampling_params_list(
+        self, req: NvCreateVideoRequest, nvext: VideoNvExt
+    ) -> tuple[list, OmniDiffusionSamplingParams]:
+        """Clone per-stage model defaults and apply explicit video overrides."""
+        defaults = list(self.engine_client.default_sampling_params_list or [])
+        if not defaults:
+            defaults = [OmniDiffusionSamplingParams()]
+            stage_types = ["diffusion"]
+        else:
+            stage_types = [
+                getattr(
+                    self.engine_client.engine.get_stage_metadata(i),
+                    "stage_type",
+                    "llm",
+                )
+                for i in range(len(defaults))
+            ]
+
+        result = []
+        output_sp = None
+        for default, stage_type in zip(defaults, stage_types, strict=True):
+            if stage_type != "diffusion":
+                result.append(
+                    default.clone() if hasattr(default, "clone") else SamplingParams()
+                )
+                continue
+
+            sp = (
+                copy.deepcopy(default)
+                if isinstance(default, OmniDiffusionSamplingParams)
+                else OmniDiffusionSamplingParams()
+            )
+            self._apply_video_sampling_overrides(sp, req, nvext)
+            result.append(sp)
+            output_sp = sp
+
+        if output_sp is None:
+            raise ValueError("Video generation requires a diffusion stage")
+        return result, output_sp
+
+    def _apply_video_sampling_overrides(
+        self,
+        sp: OmniDiffusionSamplingParams,
+        req: NvCreateVideoRequest,
+        nvext: VideoNvExt,
+    ) -> None:
+        """Overlay only fields explicitly supplied by a video request."""
+        if nvext.num_frames is not None and nvext.num_frames <= 0:
+            raise ValueError("nvext.num_frames must be greater than zero")
+        if nvext.fps is not None and nvext.fps <= 0:
+            raise ValueError("nvext.fps must be greater than zero")
+        if req.seconds is not None and req.seconds <= 0:
+            raise ValueError("seconds must be greater than zero")
+
+        if req.size is not None:
+            width, height = parse_size(req.size)
+            sp.width = width
+            sp.height = height
+
+        if nvext.num_frames is not None:
+            sp.num_frames = nvext.num_frames
+        elif req.seconds is not None:
+            frame_rate = (
+                float(nvext.fps) if nvext.fps is not None else sp.resolved_frame_rate
+            )
+            sp.num_frames = compute_num_frames(
+                seconds=req.seconds,
+                fps=frame_rate,
+                default_fps=int(
+                    getattr(self.config, "default_video_fps", DEFAULT_VIDEO_FPS)
+                ),
+            )
+        elif sp.num_frames == 1:
+            # vllm-omni uses 1 as the image-model sentinel. It is not a usable
+            # video default for pipelines that consume num_frames verbatim.
+            sp.num_frames = DEFAULT_VIDEO_NUM_FRAMES
+
+        if nvext.fps is not None:
+            sp.fps = nvext.fps
+            if hasattr(sp, "frame_rate"):
+                sp.frame_rate = float(nvext.fps)
+
+        self._update_if_not_none(sp, "num_inference_steps", nvext.num_inference_steps)
+        self._update_if_not_none(sp, "guidance_scale", nvext.guidance_scale)
+        self._update_if_not_none(sp, "seed", nvext.seed)
+        self._update_if_not_none(sp, "boundary_ratio", nvext.boundary_ratio)
+        self._update_if_not_none(sp, "guidance_scale_2", nvext.guidance_scale_2)
+        _apply_media_passthrough(sp, req.extra_args)
+
     def _engine_inputs_from_image(self, req: NvCreateImageRequest) -> EngineInputs:
         """Build engine inputs from an NvCreateImageRequest."""
         # req.size is a free-form client string, so it needs the same bound the
@@ -744,19 +838,8 @@ def _engine_inputs_from_video(
         Raises:
             ValueError: If the frame rate or output format is unsupported.
         """
-        width, height = parse_size(req.size)
         nvext = req.nvext or VideoNvExt()
 
-        num_frames = compute_num_frames(
-            num_frames=nvext.num_frames,
-            seconds=req.seconds,
-            fps=nvext.fps,
-            default_fps=DEFAULT_VIDEO_FPS,
-        )
-        fps = nvext.fps if nvext.fps is not None else DEFAULT_VIDEO_FPS
-        if fps <= 0:
-            raise ValueError(f"fps must be greater than zero, got {fps}")
-
         output_format = req.output_format.lower() if req.output_format else None
         if output_format not in (None, "mp4"):
             raise ValueError(
@@ -776,38 +859,32 @@ def _engine_inputs_from_video(
                 image.size[1],
             )
 
-        sp = OmniDiffusionSamplingParams(
-            height=height,
-            width=width,
-            num_frames=num_frames,
-        )
-        self._update_if_not_none(sp, "num_inference_steps", nvext.num_inference_steps)
-        self._update_if_not_none(sp, "guidance_scale", nvext.guidance_scale)
-        sp.seed = (
-            nvext.seed if nvext.seed is not None else random.randint(0, 2**32 - 1)
+        sampling_params_list, output_sp = self._build_video_sampling_params_list(
+            req, nvext
         )
-        self._update_if_not_none(sp, "boundary_ratio", nvext.boundary_ratio)
-        self._update_if_not_none(sp, "guidance_scale_2", nvext.guidance_scale_2)
-        self._update_if_not_none(sp, "fps", fps)
-        _apply_media_passthrough(sp, req.extra_args)
-
-        sampling_params_list = self._build_sampling_params_list(sp)
         lora_request = self._resolve_and_apply_lora(req.model, sampling_params_list)
 
+        model_fps = output_sp.resolved_frame_rate
+        output_fps = round(
+            model_fps
+            if model_fps is not None
+            else getattr(self.config, "default_video_fps", DEFAULT_VIDEO_FPS)
+        )
+
         logger.info(
             "Video diffusion request: prompt='%s...', size=%sx%s, frames=%s, fps=%s",
             req.prompt[:50],
-            width,
-            height,
-            num_frames,
-            fps,
+            getattr(output_sp, "width", None),
+            getattr(output_sp, "height", None),
+            getattr(output_sp, "num_frames", None),
+            model_fps,
         )
 
         return EngineInputs(
             prompt=prompt,
             sampling_params_list=sampling_params_list,
             request_type=RequestType.VIDEO_GENERATION,
-            fps=fps,
+            fps=output_fps,
             response_format=req.response_format,
             output_format=output_format,
             lora_request=lora_request,

components/src/dynamo/vllm/tests/omni/test_omni_args.py

@@ -3,6 +3,7 @@
 
 """Unit tests for OmniConfig validation and omni argument parsing."""
 
+import argparse
 import contextlib
 import dataclasses
 import logging
@@ -19,6 +20,7 @@
     from dynamo.vllm import main as vllm_main
     from dynamo.vllm.omni.args import (
         FlexibleArgumentParser,
+        OmniArgGroup,
         OmniConfig,
         OmniDiffusionKwargs,
         OmniEngineArgs,
@@ -105,6 +107,29 @@ def test_omni_config_valid_defaults():
     config.validate()
 
 
+def test_startup_lora_paths_parse_as_diffusion_options():
+    parser = argparse.ArgumentParser()
+    OmniArgGroup().add_arguments(parser)
+
+    args = parser.parse_args(
+        ["--lora-path", "/models/fast-a.safetensors", "/models/fast-b.safetensors"]
+    )
+
+    assert args.lora_path == [
+        "/models/fast-a.safetensors",
+        "/models/fast-b.safetensors",
+    ]
+
+
+def test_ulysses_a2a_permute_parses_as_parallel_option():
+    parser = argparse.ArgumentParser()
+    OmniArgGroup().add_arguments(parser)
+
+    args = parser.parse_args(["--ulysses-a2a-permute"])
+
+    assert args.ulysses_a2a_permute is True
+
+
 @pytest.mark.parametrize("fps", [0, -1, -100])
 def test_omni_config_invalid_video_fps(fps):
     config = _make_omni_config(default_video_fps=fps)