[Experimental][Perf] Eliminate redundant host copies from generated-video transport (#7355)

Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/user_guide/diffusion/device_side_video_postprocess.md

@@ -70,3 +70,42 @@ as safety checks or audio/video packaging.
 
 See [RFC #6541](https://github.com/vllm-project/vllm-omni/issues/6541) for the
 contract, lifecycle, batching rules, and migration plan.
+
+## Experimental registered-SHM video transport
+
+`video_output_transport.enable_registered_shm` is a separate, default-off option
+for CUDA request-mode workers. It copies eligible video tensors directly into
+CUDA-registered shared memory and returns borrowed CPU views, removing the
+intermediate pinned-to-SHM and SHM-to-consumer host copies. The D2H transfer still
+occurs.
+
+`video_output_transport.enable_borrowed_frames` independently enables borrowed
+RGB AVFrames in the Videos API MP4 response encoder, removing the third host
+copy into PyAV frame buffers. It also defaults to false. Eligible uint8 RGB
+frames are wrapped with `VideoFrame.from_numpy_buffer`; other dtypes/layouts use
+the existing encoder path. Color conversion, compression and muxing still occur.
+
+Enable it through the same CLI, Python or deploy-stage configuration:
+
+```bash
+--video-output-transport '{"enable_registered_shm": true, "enable_borrowed_frames": true}'
+```
+
+Registered SHM applies only to prepared typed video media (currently WAN2.2) and
+explicitly marked legacy video outputs (currently MiniMax-H3). H3 marks the video
+entry of its `(video, audio)` output; audio, trajectory data, images and unmarked
+legacy outputs keep their existing transport. Device-side postprocessing is
+independent and does not need to be enabled for registered SHM.
+Pre-encoded MP4 outputs bypass raw-frame transport and API-side encoding.
+
+The existing SHM routing threshold is unchanged: the tensor or its backing
+storage must exceed 1,000,000 bytes. This is not a measured performance crossover.
+SHM is allocated and registered on every transfer, without pooling; small outputs
+can regress and registration latency can vary. Benchmark representative output
+sizes before opting in. CPU outputs, non-CUDA backends, and synchronous/step-mode
+transfers keep their existing path.
+
+Borrowed views keep the unlinked mapping alive until their last user releases
+it. CUDA registration/copy failures fail the request; if CUDA cleanup cannot be
+confirmed, the worker retains the mapping for safety and must be restarted before
+further registered-SHM transfers.

tests/entrypoints/openai_api/test_video_server.py

@@ -151,6 +151,48 @@ def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixt
     handler.shutdown()
 
 
+@pytest.mark.parametrize("typed_stage", [False, True], ids=["legacy", "typed"])
+def test_remote_diffusion_stage_enables_borrowed_frames_for_both_video_responses(
+    mocker: MockerFixture, typed_stage: bool
+):
+    from vllm_omni.config.config_factory import StageConfigFactory
+
+    options = {
+        "model_class_name": "MiniMaxH3Pipeline",
+        "video_output_transport": {"enable_borrowed_frames": True},
+    }
+    stage = (
+        StageConfigFactory.create_typed_default_diffusion("test-model", options).stage_configs[0]
+        if typed_stage
+        else StageConfigFactory.create_default_diffusion(options)[0]
+    )
+    # Like StageDiffusionClient, this frontend engine exposes metadata but no
+    # worker od_config. The effective setting is available in stage_configs.
+    engine = FakeAsyncOmni()
+    engine.stage_configs = [stage]
+    handler = OmniOpenAIServingVideo.for_diffusion(engine, model_name="test-model", stage_configs=engine.stage_configs)
+    raw_encoder = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"encoded-video",
+    )
+    base64_encoder = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        return_value="encoded-video",
+    )
+    try:
+
+        async def generate_both():
+            request = VideoGenerationRequest(prompt="test prompt")
+            await handler.generate_video_bytes(request, "raw-request")
+            await handler.generate_videos(request, "base64-request")
+
+        asyncio.run(generate_both())
+        assert raw_encoder.call_args.kwargs["enable_borrowed_frames"] is True
+        assert base64_encoder.call_args.kwargs["enable_borrowed_frames"] is True
+    finally:
+        handler.shutdown()
+
+
 @pytest.mark.parametrize("batch_frames", [0, -1, True, 1.5, "17", None])
 def test_preencode_rejects_invalid_batch_frames_before_generation(batch_frames):
     engine = FakeAsyncOmni()

vllm_omni/diffusion/data.py

@@ -807,10 +807,16 @@ def uses_diffusers_adapter(od_config: object) -> bool:
 @dataclass
 class VideoOutputTransportConfig:
     enable_device_postprocess: bool = False
+    enable_registered_shm: bool = False
+    enable_borrowed_frames: bool = False
 
     def __post_init__(self) -> None:
         if not isinstance(self.enable_device_postprocess, bool):
             raise TypeError("enable_device_postprocess must be a bool")
+        if not isinstance(self.enable_registered_shm, bool):
+            raise TypeError("enable_registered_shm must be a bool")
+        if not isinstance(self.enable_borrowed_frames, bool):
+            raise TypeError("enable_borrowed_frames must be a bool")
 
 
 @dataclass
@@ -1826,11 +1832,16 @@ class DiffusionOutput:
     # Internal control-plane event emitted on first scheduler admission.
     request_started: bool = False
 
-    # Typed video-media contract. Declared last so the pre-existing positional
+    # Typed video-media contract. Appended so the pre-existing positional
     # constructor order (output, trajectory_timesteps, ...) that out-of-tree
     # pipelines rely on is preserved. Mutually exclusive with ``output``.
     media: DiffusionMediaOutput | None = None
 
+    # Compatibility adapter for joint (video, audio) outputs that have not
+    # migrated to typed media. Only this tuple entry is eligible for video
+    # transport optimizations; unmarked legacy outputs retain their old path.
+    video_output_index: int | None = None
+
     def __post_init__(self) -> None:
         if self.media is not None and not isinstance(self.media, DiffusionMediaOutput):
             raise TypeError(f"media must be DiffusionMediaOutput, got {type(self.media).__name__}")

vllm_omni/diffusion/ipc.py

@@ -12,10 +12,12 @@ serialised through the queue.
 from __future__ import annotations
 
 from collections.abc import Iterator
-from typing import Any
+from contextlib import suppress
+from typing import Any, TypeGuard
 
 import numpy as np
 import torch
+from vllm.logger import init_logger
 
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.media import (
@@ -33,49 +35,144 @@ from vllm_omni.diffusion.media import (
 _SHM_TENSOR_THRESHOLD = 1_000_000  # 1 MB
 DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result"
 _DIFFUSION_MEDIA_WIRE_TYPE = "diffusion_media_v1"
+logger = init_logger(__name__)
+
+# An uncertain CUDA cleanup must not unmap memory still used by the device.
+# Retain failed transfers until worker exit and reject further registrations.
+_failed_shm_copies: list[tuple[Any, torch.Tensor, torch.Tensor, torch.Stream]] = []
 
 # Sentinel so compute-then-assign packing can tell "field not packed" apart from
 # "field packed to None".
 _UNSET = object()
 
 
-def _array_to_shm(array: np.ndarray) -> dict[str, Any]:
-    """Copy a contiguous NumPy-compatible array into shared memory."""
+class _SharedMemoryLease:
+    """Keep an unlinked shared-memory mapping alive while array views use it."""
+
+    def __init__(self, shm: Any, array: np.ndarray) -> None:
+        self.shm = shm
+        self.__array_interface__ = array.__array_interface__
+
+    def __del__(self) -> None:
+        with suppress(BufferError):
+            self.shm.close()
+
+
+def _new_shm_array(shape: tuple[int, ...], dtype: np.dtype) -> tuple[Any, np.ndarray, dict[str, Any]]:
     from multiprocessing import shared_memory
 
+    nbytes = int(np.prod(shape, dtype=np.int64)) * dtype.itemsize
+    shm = shared_memory.SharedMemory(create=True, size=nbytes)
+    array = np.ndarray(shape, dtype=dtype, buffer=shm.buf[:nbytes])
+    handle = {
+        "name": shm.name,
+        "shape": list(shape),
+        "numpy_dtype": str(dtype),
+        "nbytes": nbytes,
+    }
+    return shm, array, handle
+
+
+def _array_to_shm(array: np.ndarray) -> dict[str, Any]:
+    """Copy a contiguous NumPy-compatible array into shared memory."""
     if array.dtype.hasobject:
         raise TypeError("NumPy object arrays cannot be transferred through raw shared memory")
 
     array = np.ascontiguousarray(array)
-    nbytes = array.nbytes
-    shm = shared_memory.SharedMemory(create=True, size=nbytes)
-    shm_array = np.ndarray(array.shape, dtype=array.dtype, buffer=shm.buf[:nbytes])
+    shm, shm_array, handle = _new_shm_array(array.shape, array.dtype)
     np.copyto(shm_array, array)
-    handle = {
-        "name": shm.name,
-        "shape": list(array.shape),
-        "numpy_dtype": str(array.dtype),
-        "nbytes": nbytes,
-    }
+    del shm_array
     shm.close()
     return handle
 
 
 def _array_from_shm(handle: dict[str, Any]) -> np.ndarray:
-    """Copy an array from shared memory, then close and unlink its segment."""
+    """Materialize or borrow an array from shared memory and unlink its name."""
     from multiprocessing import shared_memory
 
     shm = shared_memory.SharedMemory(name=handle["name"])
-    try:
     array = np.ndarray(
         handle["shape"],
         dtype=np.dtype(handle["numpy_dtype"]),
         buffer=shm.buf[: handle["nbytes"]],
-        ).copy()
+    )
+    if handle.get("borrow_on_unpack"):
+        # A base owner survives np.asarray(), slicing, and torch.from_numpy().
+        # An ndarray subclass attribute does not survive np.asarray().
+        borrowed = np.asarray(_SharedMemoryLease(shm, array))
+        shm.unlink()
+        return borrowed
+    try:
+        return array.copy()
     finally:
+        del array
         shm.close()
         shm.unlink()
-    return array
+
+
+def _cuda_registered_shm_copy(tensor: torch.Tensor, d2h_stream: torch.Stream) -> dict[str, Any]:
+    """Copy a CUDA tensor directly into a CUDA-registered POSIX SHM segment."""
+    if _failed_shm_copies:
+        raise RuntimeError("A previous registered SHM transfer failed CUDA cleanup; restart the worker")
+    original_dtype = tensor.dtype
+    source = tensor.detach()
+    # NumPy has no BF16 dtype. Transfer the original bits, without allocating
+    # an FP32 accelerator tensor or converting the borrowed CPU result.
+    if original_dtype == torch.bfloat16:
+        numpy_dtype = np.dtype(np.uint16)
+    else:
+        numpy_dtype = torch.empty((), dtype=source.dtype).numpy().dtype
+    cudart = torch.cuda.cudart()
+    shm, shm_array, handle = _new_shm_array(tuple(source.shape), numpy_dtype)
+    pointer = shm_array.ctypes.data
+    host_tensor = None
+    registered = False
+    succeeded = False
+    try:
+        host_tensor = torch.from_numpy(shm_array).view(original_dtype)
+        error = cudart.cudaHostRegister(pointer, handle["nbytes"], 0)
+        if int(error) != 0:
+            raise RuntimeError(f"cudaHostRegister failed: {cudart.cudaGetErrorString(error)}")
+        registered = True
+        old_stream = torch.accelerator.current_stream()
+        torch.accelerator.set_stream(d2h_stream)
+        try:
+            host_tensor.copy_(source, non_blocking=True)
+        finally:
+            torch.accelerator.set_stream(old_stream)
+        d2h_stream.synchronize()
+        error = cudart.cudaHostUnregister(pointer)
+        if int(error) != 0:
+            raise RuntimeError(f"cudaHostUnregister failed: {cudart.cudaGetErrorString(error)}")
+        registered = False
+        handle["borrow_on_unpack"] = True
+        handle.update(
+            {
+                "__tensor_shm__": True,
+                "torch_dtype": str(original_dtype),
+                "torch_dtype_view": original_dtype == torch.bfloat16,
+            }
+        )
+        succeeded = True
+        return handle
+    finally:
+        if registered:
+            try:
+                # A copy may have been enqueued before a later operation failed.
+                d2h_stream.synchronize()
+                error = cudart.cudaHostUnregister(pointer)
+                if int(error) != 0:
+                    raise RuntimeError(f"cudaHostUnregister failed: {cudart.cudaGetErrorString(error)}")
+                registered = False
+            except Exception:
+                assert host_tensor is not None
+                _failed_shm_copies.append((shm, host_tensor, source, d2h_stream))
+                logger.exception("Retaining registered SHM after CUDA cleanup failure; worker restart required")
+        del host_tensor, shm_array
+        if not registered:
+            shm.close()
+        if not succeeded:
+            shm.unlink()
 
 
 def _unlink_shm_handle(handle: dict[str, Any]) -> None:
@@ -111,6 +208,8 @@ def _tensor_to_shm(
     tensor: torch.Tensor,
     d2h_stream: torch.Stream | None = None,
     created: list[dict[str, Any]] | None = None,
+    *,
+    enable_registered_shm: bool = False,
 ) -> dict[str, Any]:
     """Copy a tensor into POSIX shared memory and return a metadata handle.
 
@@ -123,6 +222,17 @@ def _tensor_to_shm(
     packed.
     """
     original_dtype = tensor.dtype
+    if (
+        d2h_stream is not None
+        and enable_registered_shm
+        and tensor.device.type == "cuda"
+        # ROCm also uses the "cuda" device type, but not this CUDA runtime API.
+        and torch.version.cuda is not None
+    ):
+        handle = _cuda_registered_shm_copy(tensor, d2h_stream)
+        if created is not None:
+            created.append(handle)
+        return handle
     if d2h_stream is not None:
         # Non-blocking D2H: copy on side stream to pinned CPU memory.
         old_stream = torch.accelerator.current_stream()
@@ -165,7 +275,7 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
     if torch_dtype_str:
         original_dtype = getattr(torch, torch_dtype_str.replace("torch.", ""), None)
         if original_dtype is not None and tensor.dtype != original_dtype:
-            tensor = tensor.to(original_dtype)
+            tensor = tensor.view(original_dtype) if handle.get("torch_dtype_view") else tensor.to(original_dtype)
     return tensor
 
 
@@ -173,6 +283,8 @@ def _pack_tensor_if_large(
     val: torch.Tensor,
     d2h_stream: torch.Stream | None = None,
     created: list[dict[str, Any]] | None = None,
+    *,
+    enable_registered_shm: bool = False,
 ) -> torch.Tensor | dict:
     """Replace a tensor with an SHM handle if it exceeds the threshold.
 
@@ -187,7 +299,7 @@ def _pack_tensor_if_large(
     except Exception:
         storage_bytes = view_bytes
     if max(view_bytes, storage_bytes) > _SHM_TENSOR_THRESHOLD:
-        return _tensor_to_shm(val, d2h_stream=d2h_stream, created=created)
+        return _tensor_to_shm(val, d2h_stream=d2h_stream, created=created, enable_registered_shm=enable_registered_shm)
     return val
 
 
@@ -255,12 +367,15 @@ def _pack_diffusion_media(
     *,
     d2h_stream: torch.Stream | None,
     created: list[dict[str, Any]] | None = None,
+    enable_registered_shm: bool = False,
 ) -> dict[str, Any]:
     media.validate()
     if not media.prepared_for_transport:
         raise ValueError("Diffusion media must be prepared before IPC packing")
     video = media.video
-    packed_tensor = _pack_tensor_if_large(video.tensor, d2h_stream=d2h_stream, created=created)
+    packed_tensor = _pack_tensor_if_large(
+        video.tensor, d2h_stream=d2h_stream, created=created, enable_registered_shm=enable_registered_shm
+    )
     if isinstance(packed_tensor, torch.Tensor) and packed_tensor.device.type != "cpu":
         packed_tensor = packed_tensor.detach().cpu().contiguous()
     return {
@@ -319,6 +434,8 @@ def _pack_diffusion_fields(
     d2h_stream: torch.Stream | None = None,
     created: list[dict[str, Any]] | None = None,
     pending_updates: list[tuple[DiffusionOutput, str, object]] | None = None,
+    *,
+    enable_registered_shm: bool = False,
 ) -> DiffusionOutput:
     if output.media is not None and output.output is not None:
         raise ValueError("DiffusionOutput cannot contain both media and legacy output")
@@ -329,11 +446,26 @@ def _pack_diffusion_fields(
     # (never enqueue a half-packed object); the caller unlinks any SHM segments
     # accumulated in ``created`` on failure.
     packed_output = _UNSET
-    if output.output is not None:
+    if enable_registered_shm and output.video_output_index is not None:
+        index = output.video_output_index
+        values = output.output
+        if not isinstance(values, tuple) or type(index) is not int or not 0 <= index < len(values):
+            raise ValueError("video_output_index must identify a video tensor in the legacy output tuple")
+        if not isinstance(values[index], torch.Tensor):
+            raise TypeError("The legacy video output must be a tensor")
+        packed_output = tuple(
+            _pack_tensor_if_large(value, d2h_stream=d2h_stream, created=created, enable_registered_shm=True)
+            if position == index
+            else _pack_value_if_large(value, d2h_stream=d2h_stream, created=created)
+            for position, value in enumerate(values)
+        )
+    elif output.output is not None:
         packed_output = _pack_value_if_large(output.output, d2h_stream=d2h_stream, created=created)
     packed_media = _UNSET
     if output.media is not None:
-        packed_media = _pack_diffusion_media(output.media, d2h_stream=d2h_stream, created=created)
+        packed_media = _pack_diffusion_media(
+            output.media, d2h_stream=d2h_stream, created=created, enable_registered_shm=enable_registered_shm
+        )
     packed_latents = _UNSET
     if isinstance(output.trajectory_latents, torch.Tensor):
         packed_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream, created=created)
@@ -363,13 +495,15 @@ def _pack_diffusion_fields(
     return output
 
 
-def _is_rpc_result_envelope(output: object) -> bool:
+def _is_rpc_result_envelope(output: object) -> TypeGuard[dict[str, Any]]:
     return isinstance(output, dict) and output.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE
 
 
 def pack_diffusion_output_shm(
     output: object,
     d2h_stream: torch.Stream | None = None,
+    *,
+    enable_registered_shm: bool = False,
 ) -> object:
     """Replace large tensors in diffusion worker outputs with SHM handles.
 
@@ -382,6 +516,10 @@ def pack_diffusion_output_shm(
     If *d2h_stream* is provided, D2H copies use that stream (non-blocking on
     the default stream).  The caller must synchronize *d2h_stream* afterward.
 
+    ``enable_registered_shm`` opts explicitly identified video tensors into
+    direct CUDA-to-SHM transfer and borrowed CPU views. Other payloads keep
+    their existing transport. Registration is per transfer, without pooling.
+
     Packing is failure-atomic: every SHM segment created during the call is
     tracked and unlinked if any part of the payload fails to pack, and field
     mutations are committed only after the complete payload succeeds.
@@ -394,6 +532,7 @@ def pack_diffusion_output_shm(
             d2h_stream=d2h_stream,
             created=created,
             pending_updates=pending_updates,
+            enable_registered_shm=enable_registered_shm,
         )
     except BaseException:
         _unlink_shm_handles(created)
@@ -444,13 +583,16 @@ def _pack_output_shm(
     d2h_stream: torch.Stream | None,
     created: list[dict[str, Any]],
     pending_updates: list[tuple[DiffusionOutput, str, object]],
+    enable_registered_shm: bool = False,
 ) -> object:
+    transport_options = {"enable_registered_shm": True} if enable_registered_shm else {}
     for item in _iter_diffusion_outputs(output):
         _pack_diffusion_fields(
             item,
             d2h_stream=d2h_stream,
             created=created,
             pending_updates=pending_updates,
+            **transport_options,
         )
     return output
 
@@ -484,7 +626,7 @@ def unpack_diffusion_output_shm(output: object) -> object:
 
     result = getattr(output, "result", None)
     if isinstance(result, DiffusionOutput):
-        output.result = _unpack_diffusion_fields(result)
+        setattr(output, "result", _unpack_diffusion_fields(result))
 
     runner_outputs = getattr(output, "runner_outputs", None)
     if isinstance(runner_outputs, list):

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -2686,6 +2686,7 @@ class MiniMaxH3Pipeline(
             audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0)
         return DiffusionOutput(
             output=(video, audio),
+            video_output_index=0 if isinstance(video, torch.Tensor) else None,
             post_process_func=get_minimax_h3_post_process_func(self.od_config),
             stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}),
         )
@@ -3049,6 +3050,7 @@ class MiniMaxH3Pipeline(
             self._release_stage_cache()
         return DiffusionOutput(
             output=(video, audio),
+            video_output_index=0 if isinstance(video, torch.Tensor) else None,
             post_process_func=get_minimax_h3_post_process_func(self.od_config),
             stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}),
         )

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -1267,6 +1267,9 @@ class WorkerProc:
         """
         device = torch.device(torch.accelerator.current_accelerator().type, self.gpu_id)
         d2h_stream = torch.Stream(device=device)
+        transport_options = {}
+        if self.od_config.video_output_transport.enable_registered_shm is True:
+            transport_options["enable_registered_shm"] = True
         while True:
             item = self._async_output_queue.get()
             if item is None:
@@ -1277,7 +1280,7 @@ class WorkerProc:
                 # writing the output tensors before the side stream reads.
                 if gpu_event is not None:
                     d2h_stream.wait_event(gpu_event)
-                pack_diffusion_output_shm(output, d2h_stream=d2h_stream)
+                pack_diffusion_output_shm(output, d2h_stream=d2h_stream, **transport_options)
                 d2h_stream.synchronize()
 
                 self._enqueue_result(

vllm_omni/diffusion/worker/request_batch.py

@@ -51,6 +51,7 @@ def split_diffusion_output_by_request(
 
     return [
         DiffusionOutput(
+            video_output_index=result.video_output_index,
             output=_slice_request_output(
                 result.output,
                 idx * num_outputs_per_prompt,

vllm_omni/entrypoints/cli/serve.py

@@ -791,7 +791,10 @@ class OmniServeCommand(CLISubcommand):
             type=_json_object,
             default=None,
             help=(
-                "JSON object configuring video output preparation, for example '{\"enable_device_postprocess\": true}'."
+                "JSON object configuring video output transport: enable_device_postprocess enables GPU preparation; "
+                "enable_registered_shm enables experimental CUDA-registered SHM for supported video outputs. "
+                "enable_borrowed_frames enables borrowed RGB frames in MP4 response encoding. "
+                "All default to false."
             ),
         )
         omni_config_group.add_argument(

vllm_omni/entrypoints/openai/serving_video.py

@@ -33,7 +33,7 @@ from vllm_omni.entrypoints.openai.stage_params import (
     build_stage_sampling_params_list,
     get_default_sampling_params_list,
 )
-from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline, parse_lora_request
+from vllm_omni.entrypoints.openai.utils import get_stage_type, is_video_generation_pipeline, parse_lora_request
 from vllm_omni.entrypoints.openai.video_api_utils import (
     _encode_video_bytes,
     _PlanarFrameConverter,
@@ -187,6 +187,24 @@ class OmniOpenAIServingVideo:
         model_archs.extend(_stage_diffusion_model_class_name(stage) for stage in self.stage_configs or ())
         return od_config, tuple(get_diffusion_model_metadata(model_arch) for model_arch in model_archs)
 
+    def _video_encoding_options(self) -> dict[str, bool]:
+        transport = _config_value(self._resolve_diffusion_od_config(), "video_output_transport")
+        if transport is None:
+            # Remote diffusion clients expose only model metadata to the API.
+            # The resolved stage config still carries the effective transport options.
+            stages = self.stage_configs or getattr(self._engine_client, "stage_configs", None) or ()
+            for stage in stages:
+                if get_stage_type(stage) != "diffusion":
+                    continue
+                transport = _config_value(_config_value(stage, "diffusion_config"), "video_output_transport")
+                if transport is None:
+                    transport = _config_value(_config_value(stage, "engine_args"), "video_output_transport")
+                if transport is not None:
+                    break
+        if _config_value(transport, "enable_borrowed_frames", False) is True:
+            return {"enable_borrowed_frames": True}
+        return {}
+
     def _resolve_video_generation_defaults(
         self,
         request: VideoGenerationRequest,
@@ -535,6 +553,8 @@ class OmniOpenAIServingVideo:
             if "video_codec_options" in request.extra_params:
                 video_codec_options = request.extra_params["video_codec_options"]
 
+        encoding_options = self._video_encoding_options()
+
         def encode_video_result(idx: int, video: Any) -> str:
             if isinstance(video, bytes):
                 return base64.b64encode(video).decode("utf-8")
@@ -545,6 +565,7 @@ class OmniOpenAIServingVideo:
                 audio_sample_rate=artifacts.audio_sample_rate,
                 video_codec_options=video_codec_options,
                 frame_converter=self._video_frame_converter,
+                **encoding_options,
             )
 
         _t_encode_start = time.perf_counter()
@@ -604,6 +625,9 @@ class OmniOpenAIServingVideo:
             logger.info("Action-only video request %s completed; skipping MP4 encoding.", reference_id)
             return b"", artifacts.stage_durations, artifacts.peak_memory_mb, action, video_metadata
 
+        encoding_options: dict[str, Any] = self._video_encoding_options()
+        if audio is not None:
+            encoding_options.update(audio=audio, audio_sample_rate=artifacts.audio_sample_rate)
         _t_encode_start = time.perf_counter()
         if isinstance(artifacts.videos[0], bytes):
             video_bytes = artifacts.videos[0]
@@ -619,9 +643,9 @@ class OmniOpenAIServingVideo:
         video_bytes = _encode_video_bytes(
             artifacts.videos[0],
             fps=artifacts.output_fps,
-            **({"audio": audio, "audio_sample_rate": artifacts.audio_sample_rate} if audio is not None else {}),
             video_codec_options=video_codec_options,
             frame_converter=self._video_frame_converter,
+            **encoding_options,
         )
         _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000
         logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms)

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -749,13 +749,21 @@ def _iter_planar_video_frames(
         converter.shutdown()
 
 
+def _iter_borrowed_rgb_video_frames(frames: list[np.ndarray]) -> Generator[av.VideoFrame, None, None]:
+    """Wrap RGB arrays in AVFrames without copying their pixel storage."""
+    import av
+
+    for frame in frames:
+        yield av.VideoFrame.from_numpy_buffer(frame, format="rgb24")
+
+
 def _log_video_encoding_path(
     *,
     selected_path: str,
     frames: list[np.ndarray],
     frame_shape: tuple[int, ...],
     common_dtype: np.dtype,
-    fps: int,
+    fps: int | float,
     audio: AudioInput | None,
     audio_sample_rate: int | None,
     effective_frame_conversion_workers: int,
@@ -790,7 +798,7 @@ def _encode_prepared_video_bytes_legacy(
     frames: list[np.ndarray],
     frame_shape: tuple[int, ...],
     common_dtype: np.dtype,
-    fps: int,
+    fps: int | float,
     audio: Any | None = None,
     audio_sample_rate: int | None = None,
     video_codec_options: dict[str, str] | None = None,
@@ -835,14 +843,42 @@ def _encode_video_bytes(
     audio_sample_rate: int | None = None,
     video_codec_options: dict[str, str] | None = None,
     frame_converter: _PlanarFrameConverter | None = None,
+    enable_borrowed_frames: bool = False,
 ) -> bytes:
-    """Encode a video payload through the direct planar or legacy path."""
+    """Encode through the opt-in borrowed RGB, direct planar or legacy path."""
     from vllm_omni.diffusion.utils.media_utils import mux_av_video_audio_bytes
 
     # Prepare once so validation is shared by both paths and malformed common
     # input is reported before any muxer is opened.
     frames, frame_shape, common_dtype = _prepare_video_frames(video)
     effective_audio_sample_rate = _resolve_audio_sample_rate(audio, audio_sample_rate) if audio is not None else None
+    if (
+        enable_borrowed_frames
+        and common_dtype == np.dtype(np.uint8)
+        and len(frame_shape) == 3
+        and frame_shape[-1] == 3
+        and all(frame.dtype == np.uint8 and frame.strides[1:] == (3, 1) for frame in frames)
+    ):
+        _log_video_encoding_path(
+            selected_path="borrowed_rgb",
+            frames=frames,
+            frame_shape=frame_shape,
+            common_dtype=common_dtype,
+            fps=fps,
+            audio=audio,
+            audio_sample_rate=effective_audio_sample_rate,
+            effective_frame_conversion_workers=0,
+        )
+        audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None
+        return mux_av_video_audio_bytes(
+            _iter_borrowed_rgb_video_frames(frames),
+            width=frame_shape[1],
+            height=frame_shape[0],
+            audio_waveform=audio_np,
+            fps=float(fps),
+            audio_sample_rate=effective_audio_sample_rate,
+            video_codec_options=video_codec_options,
+        )
     fallback_reason = _direct_planar_fallback_reason(
         frames,
         frame_shape,
@@ -957,6 +993,7 @@ def encode_video_base64(
     audio_sample_rate: int | None = None,
     video_codec_options: dict[str, str] | None = None,
     frame_converter: _PlanarFrameConverter | None = None,
+    enable_borrowed_frames: bool = False,
 ) -> str:
     """Encode a video (frames/array/tensor) to base64 MP4."""
     video_bytes = _encode_video_bytes(
@@ -966,5 +1003,6 @@ def encode_video_base64(
         audio_sample_rate=audio_sample_rate,
         video_codec_options=video_codec_options,
         frame_converter=frame_converter,
+        **({"enable_borrowed_frames": True} if enable_borrowed_frames else {}),
     )
     return base64.b64encode(video_bytes).decode("utf-8")