From c7abb53e24365f5ca32775823a9d885c5c344154 Mon Sep 17 00:00:00 2001
From: Johan de Ruiter <johan@deepinfra.com>
Date: Thu, 23 Jul 2026 16:29:07 +0000
Subject: [PATCH 01/10] trtllm diffusion: add guidance_scale_2 / boundary_ratio
 to DiffusionConfig

Wan 2.2 A14B is a two-expert MoE that needs dual guidance (e.g. 4.0/3.0 with a
0.875 timestep boundary). The diffusion worker only exposed a single
--default-guidance-scale; guidance_scale_2 / boundary_ratio (Wan extra_param_specs)
defaulted to None => single guidance. Add:

- DiffusionConfig.default_guidance_scale_2 / default_boundary_ratio + the CLI flags
  --default-guidance-scale-2 / --default-boundary-ratio (backend_args).
- DiffusionEngine.generate() injects them into req.params.extra_params BEFORE the
  spec-default merge so they take effect.
- video_handler forwards the config defaults (config-only; the /v1/videos nvext
  doesn't carry them).

Both default to None, so existing models (Flux, Wan single-guidance) are unaffected.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 components/src/dynamo/trtllm/backend_args.py  | 20 +++++++++++++++++++
 .../dynamo/trtllm/configs/diffusion_config.py |  5 +++++
 .../dynamo/trtllm/engines/diffusion_engine.py | 14 +++++++++++++
 .../diffusion/video_handler.py                |  6 ++++++
 4 files changed, 45 insertions(+)

diff --git a/components/src/dynamo/trtllm/backend_args.py b/components/src/dynamo/trtllm/backend_args.py
index 8c345e78dba7..d4cdbc31003c 100644
--- a/components/src/dynamo/trtllm/backend_args.py
+++ b/components/src/dynamo/trtllm/backend_args.py
@@ -449,6 +449,24 @@ def _add_diffusion_request_arguments(self, parser: argparse.ArgumentParser) -> N
             arg_type=float,
             help="Default CFG guidance scale.",
         )
+        add_argument(
+            diffusion_request_group,
+            flag_name="--default-guidance-scale-2",
+            env_var="DYN_TRTLLM_DEFAULT_GUIDANCE_SCALE_2",
+            default=None,
+            arg_type=float,
+            help="Default second-stage CFG guidance scale (Wan 2.2 MoE low-noise "
+                 "expert). None = use the pipeline default (single guidance).",
+        )
+        add_argument(
+            diffusion_request_group,
+            flag_name="--default-boundary-ratio",
+            env_var="DYN_TRTLLM_DEFAULT_BOUNDARY_RATIO",
+            default=None,
+            arg_type=float,
+            help="Default timestep boundary ratio for switching guidance scales "
+                 "(Wan 2.2). None = use the pipeline/model default.",
+        )
         # Video specific args
         add_argument(
             diffusion_request_group,
@@ -506,6 +524,8 @@ class DynamoTrtllmConfig(ConfigBase):
     default_num_images_per_prompt: int
     default_num_inference_steps: int
     default_guidance_scale: float
+    default_guidance_scale_2: Optional[float] = None
+    default_boundary_ratio: Optional[float] = None
     torch_dtype: str
     revision: Optional[str] = None
     enable_teacache: bool
diff --git a/components/src/dynamo/trtllm/configs/diffusion_config.py b/components/src/dynamo/trtllm/configs/diffusion_config.py
index ffabe3f109ae..5357a3e34cd0 100644
--- a/components/src/dynamo/trtllm/configs/diffusion_config.py
+++ b/components/src/dynamo/trtllm/configs/diffusion_config.py
@@ -70,6 +70,11 @@ class DiffusionConfig:
     default_seconds: int = 4  # Default video duration when only fps is specified
     default_num_inference_steps: int = 50
     default_guidance_scale: float = 5.0
+    # Second-stage guidance (Wan 2.2 MoE low-noise expert) + the timestep boundary at
+    # which guidance switches from stage-1 to stage-2. None → pipeline/model default
+    # (i.e. single guidance). Set both to run Wan 2.2 dual guidance (e.g. 4.0 / 3.0 @ 0.875).
+    default_guidance_scale_2: Optional[float] = None
+    default_boundary_ratio: Optional[float] = None
 
     # ── Pipeline optimization config (maps to PipelineConfig) ──
     disable_torch_compile: bool = False
diff --git a/components/src/dynamo/trtllm/engines/diffusion_engine.py b/components/src/dynamo/trtllm/engines/diffusion_engine.py
index fdacb0ebe09c..d453185d9cdc 100644
--- a/components/src/dynamo/trtllm/engines/diffusion_engine.py
+++ b/components/src/dynamo/trtllm/engines/diffusion_engine.py
@@ -224,6 +224,8 @@ def generate(
         num_images_per_prompt: int = 1,
         num_inference_steps: int = 50,
         guidance_scale: float = 5.0,
+        guidance_scale_2: Optional[float] = None,
+        boundary_ratio: Optional[float] = None,
         seed: Optional[int] = None,
     ) -> "MediaOutput":
         """Generate video/image frames from text prompt.
@@ -286,6 +288,18 @@ def generate(
             params=params,
         )
 
+        # Wan 2.2 dual-guidance overrides. guidance_scale_2 / boundary_ratio are
+        # extra_param_specs that default to None (single guidance); set them explicitly
+        # here — before the spec-default merge below (setdefault won't clobber a value
+        # already present) — so goff (e.g. 4.0/3.0 @ boundary 0.875) actually takes effect.
+        if guidance_scale_2 is not None or boundary_ratio is not None:
+            if req.params.extra_params is None:
+                req.params.extra_params = {}
+            if guidance_scale_2 is not None:
+                req.params.extra_params["guidance_scale_2"] = guidance_scale_2
+            if boundary_ratio is not None:
+                req.params.extra_params["boundary_ratio"] = boundary_ratio
+
         # Replicate the TRTLLM's visual_gen executor's _merge_defaults: fill None fields in
         # req.params with pipeline-specific defaults (universal + extra_param
         # specs), since we call pipeline.infer() directly instead of going
diff --git a/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py b/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py
index d1f73846c964..541cd3bd0c52 100644
--- a/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py
+++ b/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py
@@ -205,6 +205,10 @@ async def generate(
                 if nvext.guidance_scale is not None
                 else self.config.default_guidance_scale
             )
+            # Dual-guidance / boundary are config-only (not carried in the /v1/videos
+            # nvext); None → engine/pipeline default (single guidance).
+            guidance_scale_2 = self.config.default_guidance_scale_2
+            boundary_ratio = self.config.default_boundary_ratio
 
             logger.info(
                 f"Request {request_id}: prompt='{req.prompt[:50]}...', "
@@ -227,6 +231,8 @@ async def generate(
                     num_frames=num_frames,
                     num_inference_steps=num_inference_steps,
                     guidance_scale=guidance_scale,
+                    guidance_scale_2=guidance_scale_2,
+                    boundary_ratio=boundary_ratio,
                     seed=nvext.seed,
                 )
 

From bbac1c1f4970cb56637e760ad4dba89b5df823f1 Mon Sep 17 00:00:00 2001
From: Johan de Ruiter <johan@deepinfra.com>
Date: Fri, 24 Jul 2026 21:46:25 +0000
Subject: [PATCH 02/10] trtllm diffusion: default video response to inline
 b64_json

The video handler defaulted response_format to "url", which uploads the clip to
media_output_fs_url and returns a file reference. deepinfra's deepapi reads the
video from data[0].b64_json and never sends response_format, so it received no
video (ERR_MODEL "No video data"). Default to b64_json so the clip is returned
inline; callers can still pass response_format="url" explicitly.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../trtllm/request_handlers/diffusion/video_handler.py      | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py b/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py
index 541cd3bd0c52..8de0d6457dd3 100644
--- a/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py
+++ b/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py
@@ -240,7 +240,11 @@ async def generate(
                 raise RuntimeError("Pipeline returned no output (MediaOutput is None)")
 
             # Determine output format
-            response_format = req.response_format or "url"
+            # Default to inline b64_json: deepinfra's deepapi consumer reads the video from
+            # data[0].b64_json and never sends response_format, so a "url" default would return
+            # a file reference it can't read (-> ERR_MODEL "No video data"). Callers that want a
+            # hosted URL can still pass response_format="url" explicitly.
+            response_format = req.response_format or "b64_json"
             if response_format not in ("url", "b64_json"):
                 raise ValueError(
                     f"Unsupported response_format: {response_format!r}; expected 'url' or 'b64_json'"

From fb342e4e2a8f08ca8015418a5bd70a49a73cb3fc Mon Sep 17 00:00:00 2001
From: Johan de Ruiter <johan@deepinfra.com>
Date: Sat, 25 Jul 2026 03:57:51 +0000
Subject: [PATCH 03/10] trtllm diffusion tests: fix default-response-format
 test for b64 + cover guidance forwarding

The b64_json default change flipped VideoGenerationHandler's default from "url"
to "b64_json", which broke test_default_response_format_is_url (it asserted
upload_to_fs was called and data[0].url was set). Rewrote it to
test_default_response_format_is_b64_json (no upload; data[0].b64_json set) so the
new default is pinned.

Also added test_guidance_defaults_forwarded_to_engine: the handler must forward
config.default_guidance_scale_2 / default_boundary_ratio to engine.generate()
(dual guidance isn't carried in the /v1/videos nvext) -- the core of this PR had
no test.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../tests/test_trtllm_video_diffusion.py      | 49 +++++++++++++++++--
 1 file changed, 44 insertions(+), 5 deletions(-)

diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py b/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py
index 21b8e50fd25f..18a3f5b54c9c 100644
--- a/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py
+++ b/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py
@@ -846,8 +846,10 @@ async def test_b64_response_format(self):
         assert decoded == b"fake_mp4_bytes"
 
     @pytest.mark.asyncio
-    async def test_default_response_format_is_url(self):
-        """Test that generate() defaults to url response format."""
+    async def test_default_response_format_is_b64_json(self):
+        """generate() defaults to b64_json: deepinfra's deepapi reads the clip from
+        data[0].b64_json and never sends response_format, so a "url" default returned a
+        reference it couldn't read (-> ERR_MODEL "No video data")."""
         handler = self._make_handler()
 
         request = {
@@ -868,9 +870,46 @@ async def test_default_response_format_is_url(self):
                 results.append(result)
 
         assert len(results) == 1
-        # Default should be "url" format, so upload_to_fs should be called
-        mock_upload.assert_called_once()
-        assert results[0]["data"][0]["url"] is not None
+        # Default is b64_json now: no upload, clip returned inline.
+        mock_upload.assert_not_called()
+        assert results[0]["data"][0]["b64_json"] is not None
+        assert results[0]["data"][0].get("url") is None
+
+    @pytest.mark.asyncio
+    async def test_guidance_defaults_forwarded_to_engine(self):
+        """The handler forwards config default_guidance_scale_2 / default_boundary_ratio
+        to engine.generate() -- dual guidance isn't carried in the /v1/videos nvext."""
+        from dynamo.trtllm.request_handlers.diffusion.video_handler import (
+            VideoGenerationHandler,
+        )
+
+        mock_output = SimpleNamespace(
+            video=torch.zeros((1, 4, 64, 64, 3), dtype=torch.uint8), image=None, audio=None,
+        )
+        mock_engine = MagicMock()
+        mock_engine.generate = MagicMock(return_value=mock_output)
+        config = DiffusionConfig(
+            media_output_fs_url="file:///tmp/test_media",
+            default_guidance_scale_2=3.0,
+            default_boundary_ratio=0.875,
+        )
+        with patch(
+            "dynamo.trtllm.request_handlers.diffusion.video_handler.get_fs",
+            return_value=MagicMock(),
+        ):
+            handler = VideoGenerationHandler(engine=mock_engine, config=config)
+
+        with patch(
+            "dynamo.trtllm.request_handlers.diffusion.video_handler.encode_to_video_bytes",
+            return_value=b"fake_mp4",
+        ):
+            async for _ in handler.generate(
+                {"prompt": "p", "model": "m", "response_format": "b64_json"}, MagicMock()
+            ):
+                pass
+
+        assert mock_engine.generate.call_args.kwargs["guidance_scale_2"] == 3.0
+        assert mock_engine.generate.call_args.kwargs["boundary_ratio"] == 0.875
 
     @pytest.mark.asyncio
     async def test_error_response_on_failure(self):

From 416c11ee7ea256e904a11fc6e96487f1c5d9307a Mon Sep 17 00:00:00 2001
From: Johan de Ruiter <johan@deepinfra.com>
Date: Tue, 28 Jul 2026 00:41:05 +0000
Subject: [PATCH 04/10] trtllm diffusion: wire cache_dit backend (Wan 2.2 ~2x
 step-skip)

VisualGen supports cache_dit (DBCache/TaylorSeer/SCM) and Wan 2.2 requires it
(it rejects teacache), but the dynamo backend only ever built a TeaCacheConfig.
Add a --cache-backend selector; when "cache_dit", build a CacheDiTConfig and
pass it as VisualGenArgs.cache (defaults are tuned for few-step -> ~2x skip).
teacache path preserved via --enable-teacache / --cache-backend teacache.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 components/src/dynamo/trtllm/backend_args.py             | 9 +++++++++
 components/src/dynamo/trtllm/configs/diffusion_config.py | 7 +++++++
 components/src/dynamo/trtllm/engines/diffusion_engine.py | 9 +++++++--
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/components/src/dynamo/trtllm/backend_args.py b/components/src/dynamo/trtllm/backend_args.py
index d4cdbc31003c..b82a937baf5a 100644
--- a/components/src/dynamo/trtllm/backend_args.py
+++ b/components/src/dynamo/trtllm/backend_args.py
@@ -277,6 +277,14 @@ def _add_diffusion_arguments(self, parser: argparse.ArgumentParser) -> None:
             arg_type=float,
             help="TeaCache threshold.",
         )
+        add_argument(
+            diffusion_group,
+            flag_name="--cache-backend",
+            env_var="DYN_TRTLLM_CACHE_BACKEND",
+            default="none",
+            arg_type=str,
+            help="Step-caching backend: none, teacache, or cache_dit (Wan 2.2 supports cache_dit, not teacache).",
+        )
         add_argument(
             diffusion_group,
             flag_name="--torch-dtype",
@@ -528,6 +536,7 @@ class DynamoTrtllmConfig(ConfigBase):
     default_boundary_ratio: Optional[float] = None
     torch_dtype: str
     revision: Optional[str] = None
+    cache_backend: str
     enable_teacache: bool
     teacache_use_ret_steps: bool
     teacache_thresh: float
diff --git a/components/src/dynamo/trtllm/configs/diffusion_config.py b/components/src/dynamo/trtllm/configs/diffusion_config.py
index 5357a3e34cd0..551ff01a0d0d 100644
--- a/components/src/dynamo/trtllm/configs/diffusion_config.py
+++ b/components/src/dynamo/trtllm/configs/diffusion_config.py
@@ -102,6 +102,13 @@ class DiffusionConfig:
     # Enable dynamic weight quantization (quantize BF16 weights on-the-fly during loading)
     quant_dynamic: bool = True
 
+    # ── Cache-acceleration backend ──
+    # Step-caching backend: "none" (default), "teacache", or "cache_dit". cache_dit
+    # (DBCache/TaylorSeer/SCM) is the ~2x step-skip Wan 2.2 supports -- it REJECTS
+    # teacache -- and VisualGen's CacheDiTConfig defaults are already tuned for
+    # few-step runs, so no extra knobs are needed here.
+    cache_backend: str = "none"
+
     # ── TeaCache optimization config (maps to TeaCacheConfig) ──
     enable_teacache: bool = False
     teacache_use_ret_steps: bool = True
diff --git a/components/src/dynamo/trtllm/engines/diffusion_engine.py b/components/src/dynamo/trtllm/engines/diffusion_engine.py
index d453185d9cdc..e8f2ab3fadb1 100644
--- a/components/src/dynamo/trtllm/engines/diffusion_engine.py
+++ b/components/src/dynamo/trtllm/engines/diffusion_engine.py
@@ -159,7 +159,7 @@ def _build_diffusion_args(self) -> "VisualGenArgs":
             TorchCompileConfig,
             VisualGenArgs,
         )
-        from tensorrt_llm._torch.visual_gen.config import AttentionConfig
+        from tensorrt_llm._torch.visual_gen.config import AttentionConfig, CacheDiTConfig
 
         # Build quant_config dict if quantization is requested
         # VisualGenArgs accepts a dict in ModelOpt format and parses it via model_validator
@@ -202,7 +202,12 @@ def _build_diffusion_args(self) -> "VisualGenArgs":
         )
 
         # Add optional fields
-        if self.config.enable_teacache:
+        if self.config.cache_backend == "cache_dit":
+            # DBCache/TaylorSeer/SCM step-skip. VisualGen's CacheDiTConfig defaults
+            # are tuned for few-step runs (warmup 4, L1 thresh 0.24, <=3 continuous
+            # cached) -> ~2x skip. Wan 2.2 supports cache_dit and rejects teacache.
+            args_kwargs["cache"] = CacheDiTConfig()
+        elif self.config.enable_teacache or self.config.cache_backend == "teacache":
             args_kwargs["cache"] = TeaCacheConfig(
                 use_ret_steps=self.config.teacache_use_ret_steps,
                 teacache_thresh=self.config.teacache_thresh,

From 2a514d3584ba302e6bc785d14d5fb84e30cc4f38 Mon Sep 17 00:00:00 2001
From: Johan de Ruiter <johan@deepinfra.com>
Date: Tue, 28 Jul 2026 17:02:31 +0000
Subject: [PATCH 05/10] trtllm diffusion: software (PyAV/libx264) mp4 encode
 for NVENC-less GPUs

encode_to_video_bytes drove h264_nvenc (hardware), which has no capable
device on datacenter compute GPUs (B200/H100/A100) -> every /v1/videos
request failed at encode ("No capable devices found"). The in-tree imageio
ffmpeg has no software h264 fallback either (nvenc + libvpx-vp9 only).

Encode via PyAV/libx264 (self-contained, bundled) to a temp file instead:
works on B200 (CPU software), ~0.7s for 720p/81f, BT.709-tagged, and no
in-memory ffmpeg output pipe to stall on. Defensive batch-dim squeeze for
MediaOutput.video's (B,T,H,W,C) shape. Worker image adds `pip install av`.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../src/dynamo/common/utils/video_utils.py    | 94 +++++++++++--------
 1 file changed, 57 insertions(+), 37 deletions(-)

diff --git a/components/src/dynamo/common/utils/video_utils.py b/components/src/dynamo/common/utils/video_utils.py
index 37326d3280bc..688e663c1bb6 100644
--- a/components/src/dynamo/common/utils/video_utils.py
+++ b/components/src/dynamo/common/utils/video_utils.py
@@ -195,48 +195,68 @@ def encode_to_video_bytes(
         ImportError: If imageio is not available.
         RuntimeError: If encoding fails.
     """
+    import tempfile
+
+    import av
+
+    # Defensive squeeze: MediaOutput.video is (B, T, H, W, C) since TRT-LLM rc9.
+    if frames.ndim == 5 and frames.shape[0] == 1:
+        frames = frames[0]
+    frames = np.ascontiguousarray(frames, dtype=np.uint8)
+    num_frames, height, width, _ = frames.shape
+
+    if output_format == "mp4":
+        codec = "libx264"
+    elif output_format == "webm":
+        codec = "libvpx-vp9"
+    else:
+        raise ValueError(f"No codec specified for response format: {output_format}")
+
+    logger.info(
+        f"Encoding {num_frames} frames to {output_format} ({codec}, software) at {fps} fps"
+    )
+
+    # Software encode via PyAV: the worker runs on NVENC-less datacenter GPUs
+    # (B200/H100/A100) where h264_nvenc has no capable device, and the in-tree
+    # imageio ffmpeg has no software h264. PyAV bundles libx264. Encode to a temp
+    # file (not an in-memory pipe) for reliability, then read the bytes back.
+    tmp = tempfile.NamedTemporaryFile(suffix=f".{output_format}", delete=False)
+    tmp.close()
     try:
-        import imageio.v3 as iio
-    except ImportError:
+        container = av.open(tmp.name, mode="w")
         try:
-            import imageio as iio  # type: ignore[no-redef]
-        except ImportError:
-            raise ImportError(
-                "imageio is required for video encoding. "
-                "Install with: pip install imageio[ffmpeg]"
-            )
-
-    logger.info(f"Encoding {len(frames)} frames to {output_format} bytes at {fps} fps")
-
-    try:
-        buffer = io.BytesIO()
-
-        kwargs: dict = {"fps": fps}
-        if output_format == "webm":
-            kwargs["codec"] = "libvpx-vp9"
-        elif output_format == "mp4":
-            kwargs["codec"] = "libx264"
-        else:
-            raise ValueError(f"No codec specified for response format: {output_format}")
-
-        if hasattr(iio, "imwrite"):
-            # v3 API
-            iio.imwrite(buffer, frames, extension=f".{output_format}", **kwargs)
-        else:
-            # v2 API
-            writer = iio.get_writer(  # type: ignore[attr-defined]
-                buffer, format="FFMPEG", mode="I", **kwargs
-            )
+            stream = container.add_stream(codec, rate=fps)
+            stream.width = width
+            stream.height = height
+            stream.pix_fmt = "yuv420p"
+            if codec == "libx264":
+                stream.options = {"crf": "18", "preset": "veryfast"}
+            # BT.709 color tags so players don't render the clip washed-out.
             try:
-                for frame in frames:
-                    writer.append_data(frame)
-            finally:
-                writer.close()
-
-        video_bytes = buffer.getvalue()
+                cc = stream.codec_context
+                cc.color_primaries = 1  # AVCOL_PRI_BT709
+                cc.color_trc = 1        # AVCOL_TRC_BT709
+                cc.colorspace = 1       # AVCOL_SPC_BT709
+                cc.color_range = 1      # AVCOL_RANGE_MPEG (limited / "tv")
+            except Exception:  # noqa: BLE001
+                logger.warning("could not set BT.709 color tags on the stream")
+            for i in range(num_frames):
+                frame = av.VideoFrame.from_ndarray(frames[i], format="rgb24")
+                for packet in stream.encode(frame):
+                    container.mux(packet)
+            for packet in stream.encode():  # flush encoder
+                container.mux(packet)
+        finally:
+            container.close()
+        with open(tmp.name, "rb") as fh:
+            video_bytes = fh.read()
         logger.info(f"Encoded video to {len(video_bytes)} bytes")
         return video_bytes
-
     except Exception as e:
         logger.error(f"Failed to encode video to bytes: {e}")
         raise RuntimeError(f"Video encoding to bytes failed: {e}") from e
+    finally:
+        try:
+            os.unlink(tmp.name)
+        except OSError:
+            pass

From 1b4c824929281e4435b43b92247a6a87b0e36940 Mon Sep 17 00:00:00 2001
From: Johan de Ruiter <johan@deepinfra.com>
Date: Tue, 28 Jul 2026 21:43:46 +0000
Subject: [PATCH 06/10] trtllm diffusion: cap video-worker CPU threads (good
 neighbor on shared GPU nodes)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Video workers left uncapped size torch's intra-op pool to the physical core count
(~112 on a 224-thread box), libx264 to every core, and the Inductor compile pool
to 32 — all of which hog a shared GPU node and thrash under contention (a 720p/81f
encode took 244s vs ~1s). Add one shared helper in dynamo.common with three
independent, env-overridable knobs, called once at each video worker's startup:

  - DI_VIDEO_TORCH_THREADS   (default 12) torch.set_num_threads; gen maxes useful
                             CPU parallelism at ~12 (12 == 32 threads -> same 268s),
                             so >12 buys no speed. 12 is well under the per-GPU fair
                             share (224/8 = 28), so full speed AND neighborly.
  - DI_VIDEO_ENCODE_THREADS  (default 12) libx264 thread_count; measured PyAV knee.
  - DI_VIDEO_COMPILE_THREADS (default 32) TORCHINDUCTOR_COMPILE_THREADS (boot-only).

Also sets OMP_WAIT_POLICY=PASSIVE / KMP_BLOCKTIME=0 so idle pools sleep between the
sequential compile -> generate -> encode phases and hand their cores to the active
phase instead of spinning. The three phases never overlap, so the budgets don't sum.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../src/dynamo/common/utils/video_utils.py    | 56 +++++++++++++++++++
 .../trtllm/workers/video_diffusion_worker.py  |  6 ++
 2 files changed, 62 insertions(+)

diff --git a/components/src/dynamo/common/utils/video_utils.py b/components/src/dynamo/common/utils/video_utils.py
index 688e663c1bb6..e0dc5167336b 100644
--- a/components/src/dynamo/common/utils/video_utils.py
+++ b/components/src/dynamo/common/utils/video_utils.py
@@ -175,6 +175,61 @@ def encode_to_mp4(
         raise RuntimeError(f"Video encoding failed: {e}") from e
 
 
+# Video workers have three distinct CPU consumers that peak in different phases and want
+# different thread counts, so each gets its own env-overridable knob (defaults below):
+#   - Inductor compile pool (boot / first request: parallel Triton kernel compilation)
+#   - torch intra-op pool   (generation: CPU-side ops around a GPU-bound diffusion)
+#   - libx264 encode        (encode: ~12 is the measured PyAV sweet spot for 720p)
+# On shared GPU nodes an uncapped worker is a noisy neighbor: torch sizes its pool to the
+# physical core count (~112 on a 224-thread box) and libx264 grabs every core, which both
+# thrashes (a 720p/81f encode took 244s vs ~1s) and starves co-located pods. The three
+# phases are sequential, so with OMP_WAIT_POLICY=PASSIVE / KMP_BLOCKTIME=0 an idle pool
+# sleeps and hands its cores to the active phase instead of spinning — they don't sum.
+DEFAULT_VIDEO_TORCH_THREADS = 12
+DEFAULT_VIDEO_ENCODE_THREADS = 12
+DEFAULT_VIDEO_COMPILE_THREADS = 32
+
+
+def video_encode_threads() -> int:
+    """CPU threads for the libx264 software encode (env: DI_VIDEO_ENCODE_THREADS)."""
+    return int(os.getenv("DI_VIDEO_ENCODE_THREADS", str(DEFAULT_VIDEO_ENCODE_THREADS)))
+
+
+def limit_video_worker_threads() -> None:
+    """Cap a video worker's CPU threads so it stays a good neighbor on shared nodes.
+
+    Call once at the very top of a video worker's startup, before torch loads. Sets the
+    OpenMP wait policy (idle pools sleep between the sequential compile -> generate ->
+    encode phases rather than spinning), bounds the Inductor compile pool, and caps
+    torch's intra-op pool. The encoder is capped separately via video_encode_threads().
+    Every video worker entry calls this once — that is the general pattern. No-op if
+    torch is unavailable.
+    """
+    # Import-time knobs: set before torch imports so OpenMP / Inductor pick them up. Idle
+    # pools then sleep instead of spin (covers both libgomp and Intel OpenMP / MKL).
+    os.environ.setdefault("OMP_WAIT_POLICY", "PASSIVE")
+    os.environ.setdefault("KMP_BLOCKTIME", "0")
+    os.environ.setdefault(
+        "TORCHINDUCTOR_COMPILE_THREADS",
+        os.getenv("DI_VIDEO_COMPILE_THREADS", str(DEFAULT_VIDEO_COMPILE_THREADS)),
+    )
+    torch_threads = int(
+        os.getenv("DI_VIDEO_TORCH_THREADS", str(DEFAULT_VIDEO_TORCH_THREADS))
+    )
+    try:
+        import torch
+
+        torch.set_num_threads(torch_threads)
+    except Exception:  # noqa: BLE001
+        logger.warning("could not cap torch CPU threads to %s", torch_threads)
+    logger.info(
+        "video worker CPU budget: torch=%s encode=%s compile=%s (OMP_WAIT_POLICY=PASSIVE)",
+        torch_threads,
+        video_encode_threads(),
+        os.environ.get("TORCHINDUCTOR_COMPILE_THREADS"),
+    )
+
+
 def encode_to_video_bytes(
     frames: np.ndarray,
     fps: int = 16,
@@ -229,6 +284,7 @@ def encode_to_video_bytes(
             stream.width = width
             stream.height = height
             stream.pix_fmt = "yuv420p"
+            stream.codec_context.thread_count = video_encode_threads()
             if codec == "libx264":
                 stream.options = {"crf": "18", "preset": "veryfast"}
             # BT.709 color tags so players don't render the clip washed-out.
diff --git a/components/src/dynamo/trtllm/workers/video_diffusion_worker.py b/components/src/dynamo/trtllm/workers/video_diffusion_worker.py
index 417adc4207ed..dfca15a416e6 100644
--- a/components/src/dynamo/trtllm/workers/video_diffusion_worker.py
+++ b/components/src/dynamo/trtllm/workers/video_diffusion_worker.py
@@ -33,6 +33,12 @@ async def init_video_diffusion_worker(
         shutdown_event: Event to signal shutdown.
         shutdown_endpoints: Optional list to populate with endpoints for graceful shutdown.
     """
+    from dynamo.common.utils.video_utils import limit_video_worker_threads
+
+    # Good-neighbor CPU budget on shared GPU nodes (bounds torch's intra-op pool; the
+    # encoder is bounded via the same budget). Set before the engine/torch load.
+    limit_video_worker_threads()
+
     # Check tensorrt_llm visual_gen availability early with a clear error message.
     # visual_gen is part of TensorRT-LLM (tensorrt_llm._torch.visual_gen).
     # Without this check, users would get a cryptic ImportError deep inside

From 112d0d42637f46469ebea0d19c3cff2a091f36bf Mon Sep 17 00:00:00 2001
From: Johan de Ruiter <johan@deepinfra.com>
Date: Tue, 28 Jul 2026 21:56:00 +0000
Subject: [PATCH 07/10] trtllm diffusion tests: realign encode tests to PyAV +
 cover thread caps

The PyAV rewrite (51d380b8c) replaced imageio/BytesIO with av/tempfile in
encode_to_video_bytes but left the whole TestEncodeToVideoBytes class mocking
imageio.v3 + io.BytesIO -- stale against the real code (and asserting the wrong
exception type: unsupported format now raises ValueError, not RuntimeError).

Rewrite them to exercise the real PyAV encoder (pytest.importorskip("av")):
decode the output and assert h264 / vp9, BT.709 color tags, the (B,T,H,W,C)->
(T,H,W,C) squeeze, and ValueError on an unsupported container. Add coverage for
the new thread-cap helper: DI_VIDEO_ENCODE_THREADS default/override, and
limit_video_worker_threads setting OMP_WAIT_POLICY / KMP_BLOCKTIME /
TORCHINDUCTOR_COMPILE_THREADS + torch.set_num_threads (with DI_VIDEO_TORCH_THREADS
override).

Verified: 9 passed under pytest 9.0.3 in the rc14 runtime image.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../dynamo/common/tests/test_video_utils.py   | 191 ++++++++----------
 1 file changed, 87 insertions(+), 104 deletions(-)

diff --git a/components/src/dynamo/common/tests/test_video_utils.py b/components/src/dynamo/common/tests/test_video_utils.py
index fab867fb611c..7d596c9c566e 100644
--- a/components/src/dynamo/common/tests/test_video_utils.py
+++ b/components/src/dynamo/common/tests/test_video_utils.py
@@ -3,6 +3,7 @@
 
 """Unit tests for dynamo.common.utils.video_utils module."""
 
+import io
 from unittest.mock import MagicMock, patch
 
 import numpy as np
@@ -15,142 +16,124 @@
 ]
 
 
-def make_frames(n=3, h=8, w=8) -> np.ndarray:
-    """Return a small uint8 frame array (n, h, w, 3)."""
+def make_frames(n=3, h=16, w=16) -> np.ndarray:
+    """Return a small uint8 frame array (n, h, w, 3). Even dims for libx264."""
     return np.zeros((n, h, w, 3), dtype=np.uint8)
 
 
 # ---------------------------------------------------------------------------
-# encode_to_video_bytes
+# encode_to_video_bytes — PyAV/libx264 software encode (no NVENC, no imageio)
 # ---------------------------------------------------------------------------
 
 
 class TestEncodeToVideoBytes:
-    """Tests for encode_to_video_bytes()."""
-
-    def _mock_iio_v3(self):
-        """Return a mock that looks like imageio.v3 (has imwrite)."""
-        iio = MagicMock()
-        iio.imwrite = MagicMock()
-        return iio
-
-    def _mock_iio_v2(self):
-        """Return a mock that looks like imageio v2 (no imwrite, has get_writer)."""
-        iio = MagicMock(spec=[])  # no attributes by default
-        writer = MagicMock()
-        iio.get_writer = MagicMock(return_value=writer)
-        return iio, writer
-
-    def test_mp4_selects_libx264_codec(self):
-        from dynamo.common.utils.video_utils import encode_to_video_bytes
-
-        iio = self._mock_iio_v3()
-        with patch("dynamo.common.utils.video_utils.io") as mock_io, patch(
-            "imageio.v3", iio, create=True
-        ), patch.dict("sys.modules", {"imageio.v3": iio}):
-            buf = MagicMock()
-            buf.getvalue.return_value = b"fake-mp4"
-            mock_io.BytesIO.return_value = buf
+    """Tests for encode_to_video_bytes(): software-encodes via PyAV so it works on
+    NVENC-less datacenter GPUs. These exercise the real encoder (skipped if PyAV is
+    unavailable) rather than mocking, since the whole point is that the bytes decode."""
 
-            encode_to_video_bytes(make_frames(), fps=8, output_format="mp4")
+    def test_mp4_returns_h264_bytes(self):
+        av = pytest.importorskip("av")
+        from dynamo.common.utils.video_utils import encode_to_video_bytes
 
-            iio.imwrite.assert_called_once()
-            _, kwargs = iio.imwrite.call_args
-            assert kwargs.get("codec") == "libx264"
-            assert kwargs.get("fps") == 8
+        out = encode_to_video_bytes(make_frames(n=5), fps=8, output_format="mp4")
+        assert isinstance(out, bytes) and len(out) > 0
+        with av.open(io.BytesIO(out)) as container:
+            assert container.streams.video[0].codec_context.name == "h264"
 
-    def test_webm_selects_libvpx_vp9_codec(self):
+    def test_webm_returns_vp9_bytes(self):
+        av = pytest.importorskip("av")
         from dynamo.common.utils.video_utils import encode_to_video_bytes
 
-        iio = self._mock_iio_v3()
-        with patch("dynamo.common.utils.video_utils.io") as mock_io, patch(
-            "imageio.v3", iio, create=True
-        ), patch.dict("sys.modules", {"imageio.v3": iio}):
-            buf = MagicMock()
-            buf.getvalue.return_value = b"fake-webm"
-            mock_io.BytesIO.return_value = buf
+        out = encode_to_video_bytes(make_frames(n=5), fps=8, output_format="webm")
+        assert isinstance(out, bytes) and len(out) > 0
+        with av.open(io.BytesIO(out)) as container:
+            assert container.streams.video[0].codec_context.name in ("vp9", "libvpx-vp9")
 
-            encode_to_video_bytes(make_frames(), fps=16, output_format="webm")
+    def test_mp4_tags_bt709(self):
+        av = pytest.importorskip("av")
+        from dynamo.common.utils.video_utils import encode_to_video_bytes
 
-            iio.imwrite.assert_called_once()
-            _, kwargs = iio.imwrite.call_args
-            assert kwargs.get("codec") == "libvpx-vp9"
+        out = encode_to_video_bytes(make_frames(n=5), fps=16, output_format="mp4")
+        with av.open(io.BytesIO(out)) as container:
+            cc = container.streams.video[0].codec_context
+            # BT.709 primaries/transfer/colorspace so players don't render washed-out.
+            assert int(cc.color_primaries) == 1
+            assert int(cc.color_trc) == 1
+            assert int(cc.colorspace) == 1
 
-    def test_mp4_passes_extension_to_imwrite(self):
+    def test_squeezes_5d_batch_dim(self):
+        av = pytest.importorskip("av")
         from dynamo.common.utils.video_utils import encode_to_video_bytes
 
-        iio = self._mock_iio_v3()
-        with patch("dynamo.common.utils.video_utils.io") as mock_io, patch(
-            "imageio.v3", iio, create=True
-        ), patch.dict("sys.modules", {"imageio.v3": iio}):
-            buf = MagicMock()
-            buf.getvalue.return_value = b"bytes"
-            mock_io.BytesIO.return_value = buf
+        # MediaOutput.video is (B, T, H, W, C) since TRT-LLM rc9; encode must squeeze B=1.
+        frames_5d = make_frames(n=5)[None]  # (1, 5, 16, 16, 3)
+        out = encode_to_video_bytes(frames_5d, fps=8, output_format="mp4")
+        with av.open(io.BytesIO(out)) as container:
+            assert container.streams.video[0].codec_context.name == "h264"
 
-            encode_to_video_bytes(make_frames(), output_format="mp4")
+    def test_unsupported_format_raises_value_error(self):
+        pytest.importorskip("av")
+        from dynamo.common.utils.video_utils import encode_to_video_bytes
 
-            _, kwargs = iio.imwrite.call_args
-            assert kwargs.get("extension") == ".mp4"
+        with pytest.raises(ValueError):
+            encode_to_video_bytes(make_frames(), output_format="avi")
 
-    def test_webm_passes_extension_to_imwrite(self):
-        from dynamo.common.utils.video_utils import encode_to_video_bytes
 
-        iio = self._mock_iio_v3()
-        with patch("dynamo.common.utils.video_utils.io") as mock_io, patch(
-            "imageio.v3", iio, create=True
-        ), patch.dict("sys.modules", {"imageio.v3": iio}):
-            buf = MagicMock()
-            buf.getvalue.return_value = b"bytes"
-            mock_io.BytesIO.return_value = buf
+# ---------------------------------------------------------------------------
+# limit_video_worker_threads / video_encode_threads — shared video-worker caps
+# ---------------------------------------------------------------------------
 
-            encode_to_video_bytes(make_frames(), output_format="webm")
 
-            _, kwargs = iio.imwrite.call_args
-            assert kwargs.get("extension") == ".webm"
+class TestVideoWorkerThreadCaps:
+    """The three env-overridable CPU-thread knobs + the OpenMP wait policy."""
 
-    def test_unsupported_format_raises_value_error(self):
-        from dynamo.common.utils.video_utils import encode_to_video_bytes
+    def test_encode_threads_default(self, monkeypatch):
+        monkeypatch.delenv("DI_VIDEO_ENCODE_THREADS", raising=False)
+        from dynamo.common.utils.video_utils import (
+            DEFAULT_VIDEO_ENCODE_THREADS,
+            video_encode_threads,
+        )
 
-        iio = self._mock_iio_v3()
-        with patch("dynamo.common.utils.video_utils.io") as mock_io, patch(
-            "imageio.v3", iio, create=True
-        ), patch.dict("sys.modules", {"imageio.v3": iio}):
-            mock_io.BytesIO.return_value = MagicMock()
+        assert video_encode_threads() == DEFAULT_VIDEO_ENCODE_THREADS
 
-            # ValueError is wrapped into RuntimeError by the except block
-            with pytest.raises(RuntimeError, match="Video encoding to bytes failed"):
-                encode_to_video_bytes(make_frames(), output_format="avi")
+    def test_encode_threads_env_override(self, monkeypatch):
+        monkeypatch.setenv("DI_VIDEO_ENCODE_THREADS", "7")
+        from dynamo.common.utils.video_utils import video_encode_threads
 
-    def test_returns_bytes_from_buffer(self):
-        from dynamo.common.utils.video_utils import encode_to_video_bytes
+        assert video_encode_threads() == 7
 
-        expected = b"\x00\x01\x02"
-        iio = self._mock_iio_v3()
-        with patch("dynamo.common.utils.video_utils.io") as mock_io, patch(
-            "imageio.v3", iio, create=True
-        ), patch.dict("sys.modules", {"imageio.v3": iio}):
-            buf = MagicMock()
-            buf.getvalue.return_value = expected
-            mock_io.BytesIO.return_value = buf
+    def test_limit_sets_wait_policy_compile_and_torch(self, monkeypatch):
+        import os
 
-            result = encode_to_video_bytes(make_frames(), output_format="mp4")
+        for k in (
+            "OMP_WAIT_POLICY",
+            "KMP_BLOCKTIME",
+            "TORCHINDUCTOR_COMPILE_THREADS",
+            "DI_VIDEO_COMPILE_THREADS",
+            "DI_VIDEO_TORCH_THREADS",
+        ):
+            monkeypatch.delenv(k, raising=False)
+        from dynamo.common.utils import video_utils
 
-        assert result == expected
+        fake_torch = MagicMock()
+        with patch.dict("sys.modules", {"torch": fake_torch}):
+            video_utils.limit_video_worker_threads()
 
-    def test_v2_api_fallback_writes_all_frames(self):
-        """When imageio.v3.imwrite is absent, falls back to get_writer loop."""
-        from dynamo.common.utils.video_utils import encode_to_video_bytes
+        assert os.environ["OMP_WAIT_POLICY"] == "PASSIVE"
+        assert os.environ["KMP_BLOCKTIME"] == "0"
+        assert os.environ["TORCHINDUCTOR_COMPILE_THREADS"] == str(
+            video_utils.DEFAULT_VIDEO_COMPILE_THREADS
+        )
+        fake_torch.set_num_threads.assert_called_once_with(
+            video_utils.DEFAULT_VIDEO_TORCH_THREADS
+        )
 
-        iio_v2, writer = self._mock_iio_v2()
-        with patch("dynamo.common.utils.video_utils.io") as mock_io, patch(
-            "imageio.v3", iio_v2, create=True
-        ), patch.dict("sys.modules", {"imageio.v3": iio_v2}):
-            buf = MagicMock()
-            buf.getvalue.return_value = b"v2-bytes"
-            mock_io.BytesIO.return_value = buf
+    def test_limit_torch_threads_env_override(self, monkeypatch):
+        monkeypatch.setenv("DI_VIDEO_TORCH_THREADS", "9")
+        from dynamo.common.utils import video_utils
 
-            frames = make_frames(n=4)
-            encode_to_video_bytes(frames, output_format="mp4")
+        fake_torch = MagicMock()
+        with patch.dict("sys.modules", {"torch": fake_torch}):
+            video_utils.limit_video_worker_threads()
 
-            assert writer.append_data.call_count == 4
-            writer.close.assert_called_once()
+        fake_torch.set_num_threads.assert_called_once_with(9)

From e785a5e43a1f9966ac725b9e0d542089cf584863 Mon Sep 17 00:00:00 2001
From: Johan de Ruiter <johan@deepinfra.com>
Date: Tue, 28 Jul 2026 23:09:00 +0000
Subject: [PATCH 08/10] style: pre-commit formatting (drop unused io import,
 black wraps)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 components/src/dynamo/common/tests/test_video_utils.py | 5 ++++-
 components/src/dynamo/common/utils/video_utils.py      | 7 +++----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/components/src/dynamo/common/tests/test_video_utils.py b/components/src/dynamo/common/tests/test_video_utils.py
index 7d596c9c566e..0ce3a025e61e 100644
--- a/components/src/dynamo/common/tests/test_video_utils.py
+++ b/components/src/dynamo/common/tests/test_video_utils.py
@@ -47,7 +47,10 @@ def test_webm_returns_vp9_bytes(self):
         out = encode_to_video_bytes(make_frames(n=5), fps=8, output_format="webm")
         assert isinstance(out, bytes) and len(out) > 0
         with av.open(io.BytesIO(out)) as container:
-            assert container.streams.video[0].codec_context.name in ("vp9", "libvpx-vp9")
+            assert container.streams.video[0].codec_context.name in (
+                "vp9",
+                "libvpx-vp9",
+            )
 
     def test_mp4_tags_bt709(self):
         av = pytest.importorskip("av")
diff --git a/components/src/dynamo/common/utils/video_utils.py b/components/src/dynamo/common/utils/video_utils.py
index e0dc5167336b..4f95053a2f4d 100644
--- a/components/src/dynamo/common/utils/video_utils.py
+++ b/components/src/dynamo/common/utils/video_utils.py
@@ -7,7 +7,6 @@
 video frames to MP4 format.
 """
 
-import io
 import logging
 import os
 from typing import Tuple
@@ -291,9 +290,9 @@ def encode_to_video_bytes(
             try:
                 cc = stream.codec_context
                 cc.color_primaries = 1  # AVCOL_PRI_BT709
-                cc.color_trc = 1        # AVCOL_TRC_BT709
-                cc.colorspace = 1       # AVCOL_SPC_BT709
-                cc.color_range = 1      # AVCOL_RANGE_MPEG (limited / "tv")
+                cc.color_trc = 1  # AVCOL_TRC_BT709
+                cc.colorspace = 1  # AVCOL_SPC_BT709
+                cc.color_range = 1  # AVCOL_RANGE_MPEG (limited / "tv")
             except Exception:  # noqa: BLE001
                 logger.warning("could not set BT.709 color tags on the stream")
             for i in range(num_frames):

From 0f913c3a6ec325ea9fe54e625e20f75e80cacd15 Mon Sep 17 00:00:00 2001
From: Johan de Ruiter <johan@deepinfra.com>
Date: Tue, 28 Jul 2026 23:15:00 +0000
Subject: [PATCH 09/10] style: pre-commit formatting for the
 guidance/cache/test commits (black/isort)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 components/src/dynamo/trtllm/backend_args.py               | 4 ++--
 components/src/dynamo/trtllm/engines/diffusion_engine.py   | 5 ++++-
 .../src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py | 7 +++++--
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/components/src/dynamo/trtllm/backend_args.py b/components/src/dynamo/trtllm/backend_args.py
index b82a937baf5a..79a34d5dd9d2 100644
--- a/components/src/dynamo/trtllm/backend_args.py
+++ b/components/src/dynamo/trtllm/backend_args.py
@@ -464,7 +464,7 @@ def _add_diffusion_request_arguments(self, parser: argparse.ArgumentParser) -> N
             default=None,
             arg_type=float,
             help="Default second-stage CFG guidance scale (Wan 2.2 MoE low-noise "
-                 "expert). None = use the pipeline default (single guidance).",
+            "expert). None = use the pipeline default (single guidance).",
         )
         add_argument(
             diffusion_request_group,
@@ -473,7 +473,7 @@ def _add_diffusion_request_arguments(self, parser: argparse.ArgumentParser) -> N
             default=None,
             arg_type=float,
             help="Default timestep boundary ratio for switching guidance scales "
-                 "(Wan 2.2). None = use the pipeline/model default.",
+            "(Wan 2.2). None = use the pipeline/model default.",
         )
         # Video specific args
         add_argument(
diff --git a/components/src/dynamo/trtllm/engines/diffusion_engine.py b/components/src/dynamo/trtllm/engines/diffusion_engine.py
index e8f2ab3fadb1..b2bb49d20f59 100644
--- a/components/src/dynamo/trtllm/engines/diffusion_engine.py
+++ b/components/src/dynamo/trtllm/engines/diffusion_engine.py
@@ -159,7 +159,10 @@ def _build_diffusion_args(self) -> "VisualGenArgs":
             TorchCompileConfig,
             VisualGenArgs,
         )
-        from tensorrt_llm._torch.visual_gen.config import AttentionConfig, CacheDiTConfig
+        from tensorrt_llm._torch.visual_gen.config import (
+            AttentionConfig,
+            CacheDiTConfig,
+        )
 
         # Build quant_config dict if quantization is requested
         # VisualGenArgs accepts a dict in ModelOpt format and parses it via model_validator
diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py b/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py
index 18a3f5b54c9c..6a68cc82988c 100644
--- a/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py
+++ b/components/src/dynamo/trtllm/tests/test_trtllm_video_diffusion.py
@@ -884,7 +884,9 @@ async def test_guidance_defaults_forwarded_to_engine(self):
         )
 
         mock_output = SimpleNamespace(
-            video=torch.zeros((1, 4, 64, 64, 3), dtype=torch.uint8), image=None, audio=None,
+            video=torch.zeros((1, 4, 64, 64, 3), dtype=torch.uint8),
+            image=None,
+            audio=None,
         )
         mock_engine = MagicMock()
         mock_engine.generate = MagicMock(return_value=mock_output)
@@ -904,7 +906,8 @@ async def test_guidance_defaults_forwarded_to_engine(self):
             return_value=b"fake_mp4",
         ):
             async for _ in handler.generate(
-                {"prompt": "p", "model": "m", "response_format": "b64_json"}, MagicMock()
+                {"prompt": "p", "model": "m", "response_format": "b64_json"},
+                MagicMock(),
             ):
                 pass
 

From a9f16fe95b8b8dfb6ab8f067b07a9066958079dd Mon Sep 17 00:00:00 2001
From: Johan de Ruiter <johan@deepinfra.com>
Date: Tue, 28 Jul 2026 23:27:18 +0000
Subject: [PATCH 10/10] fix: address adversarial review of the video-diffusion
 changes

- Make guidance mode + cache backend observable so goff can't silently degrade
  to single guidance (the flags are config-only; a missing --default-guidance-
  scale-2 falls back to single with no error): log DUAL/SINGLE (+ values,
  cache_backend) at worker init, add guidance to the per-request handler log,
  and include guidance_scale/guidance_scale_2/boundary_ratio/cache_backend in
  DiffusionConfig.__str__.
- backend_args: constrain --cache-backend to choices [none, teacache, cache_dit]
  (a typo previously fell through to no-cache -> ~2x slower, no error).
- video_handler: replace the exact-5D assert (AssertionError on 4D, stripped
  under -O) with explicit (T,H,W,C)/(1,T,H,W,C) handling + a clear RuntimeError.
- video_utils: warn if limit_video_worker_threads() runs after torch is already
  imported (OMP_WAIT_POLICY/KMP_BLOCKTIME would not bind) + honest docstring;
  fix the stale imageio -> PyAV Raises doc.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../src/dynamo/common/utils/video_utils.py    | 27 ++++++++++++++-----
 components/src/dynamo/trtllm/backend_args.py  |  1 +
 .../dynamo/trtllm/configs/diffusion_config.py |  4 +++
 .../diffusion/video_handler.py                | 25 ++++++++++++-----
 .../trtllm/workers/video_diffusion_worker.py  | 13 +++++++++
 5 files changed, 56 insertions(+), 14 deletions(-)

diff --git a/components/src/dynamo/common/utils/video_utils.py b/components/src/dynamo/common/utils/video_utils.py
index 4f95053a2f4d..61024f295e13 100644
--- a/components/src/dynamo/common/utils/video_utils.py
+++ b/components/src/dynamo/common/utils/video_utils.py
@@ -197,13 +197,25 @@ def video_encode_threads() -> int:
 def limit_video_worker_threads() -> None:
     """Cap a video worker's CPU threads so it stays a good neighbor on shared nodes.
 
-    Call once at the very top of a video worker's startup, before torch loads. Sets the
-    OpenMP wait policy (idle pools sleep between the sequential compile -> generate ->
-    encode phases rather than spinning), bounds the Inductor compile pool, and caps
-    torch's intra-op pool. The encoder is capped separately via video_encode_threads().
-    Every video worker entry calls this once — that is the general pattern. No-op if
-    torch is unavailable.
+    Call once at the very top of a video worker's startup, before torch loads. Caps
+    torch's intra-op pool (via torch.set_num_threads, which always applies), bounds the
+    Inductor compile pool, and sets the OpenMP wait policy so idle pools sleep between
+    the sequential compile -> generate -> encode phases rather than spinning. The encoder
+    is capped separately via video_encode_threads(). Every video worker entry calls this
+    once — that is the general pattern. No-op if torch is unavailable.
+
+    Caveat: OMP_WAIT_POLICY / KMP_BLOCKTIME only bind if set before the OpenMP runtime
+    initializes (≈ first torch import), so for guaranteed effect set them in the launch
+    env. torch.set_num_threads and TORCHINDUCTOR_COMPILE_THREADS apply regardless.
     """
+    import sys
+
+    if "torch" in sys.modules:
+        logger.warning(
+            "limit_video_worker_threads() ran after torch was already imported — "
+            "OMP_WAIT_POLICY/KMP_BLOCKTIME may not take effect; set them in the launch "
+            "env for reliability (torch.set_num_threads still applies)."
+        )
     # Import-time knobs: set before torch imports so OpenMP / Inductor pick them up. Idle
     # pools then sleep instead of spin (covers both libgomp and Intel OpenMP / MKL).
     os.environ.setdefault("OMP_WAIT_POLICY", "PASSIVE")
@@ -246,7 +258,8 @@ def encode_to_video_bytes(
         Encoded video as bytes.
 
     Raises:
-        ImportError: If imageio is not available.
+        ImportError: If PyAV (``av``) is not available.
+        ValueError: If output_format has no known codec.
         RuntimeError: If encoding fails.
     """
     import tempfile
diff --git a/components/src/dynamo/trtllm/backend_args.py b/components/src/dynamo/trtllm/backend_args.py
index 79a34d5dd9d2..8558adb24c5e 100644
--- a/components/src/dynamo/trtllm/backend_args.py
+++ b/components/src/dynamo/trtllm/backend_args.py
@@ -283,6 +283,7 @@ def _add_diffusion_arguments(self, parser: argparse.ArgumentParser) -> None:
             env_var="DYN_TRTLLM_CACHE_BACKEND",
             default="none",
             arg_type=str,
+            choices=["none", "teacache", "cache_dit"],
             help="Step-caching backend: none, teacache, or cache_dit (Wan 2.2 supports cache_dit, not teacache).",
         )
         add_argument(
diff --git a/components/src/dynamo/trtllm/configs/diffusion_config.py b/components/src/dynamo/trtllm/configs/diffusion_config.py
index 551ff01a0d0d..7a1340ebd4d4 100644
--- a/components/src/dynamo/trtllm/configs/diffusion_config.py
+++ b/components/src/dynamo/trtllm/configs/diffusion_config.py
@@ -161,6 +161,10 @@ def __str__(self) -> str:
             f"default_num_frames={self.default_num_frames}, "
             f"default_num_images_per_prompt={self.default_num_images_per_prompt}, "
             f"default_num_inference_steps={self.default_num_inference_steps}, "
+            f"default_guidance_scale={self.default_guidance_scale}, "
+            f"default_guidance_scale_2={self.default_guidance_scale_2}, "
+            f"default_boundary_ratio={self.default_boundary_ratio}, "
+            f"cache_backend={self.cache_backend}, "
             f"enable_teacache={self.enable_teacache}, "
             f"attn_backend={self.attn_backend}, "
             f"quant_algo={self.quant_algo}, "
diff --git a/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py b/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py
index 8de0d6457dd3..41de45a0d52c 100644
--- a/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py
+++ b/components/src/dynamo/trtllm/request_handlers/diffusion/video_handler.py
@@ -210,9 +210,15 @@ async def generate(
             guidance_scale_2 = self.config.default_guidance_scale_2
             boundary_ratio = self.config.default_boundary_ratio
 
+            guidance_desc = (
+                f"{guidance_scale}/{guidance_scale_2}@{boundary_ratio}"
+                if guidance_scale_2 is not None
+                else f"{guidance_scale} (single)"
+            )
             logger.info(
                 f"Request {request_id}: prompt='{req.prompt[:50]}...', "
-                f"size={width}x{height}, frames={num_frames}, steps={num_inference_steps}"
+                f"size={width}x{height}, frames={num_frames}, steps={num_inference_steps}, "
+                f"guidance={guidance_desc}"
             )
 
             # Run generation in thread pool (blocking operation).
@@ -258,13 +264,18 @@ async def generate(
 
             # Encode media based on what the pipeline returned
             if output.video is not None:
-                # MediaOutput.video is (B, T, H, W, C) uint8 since TRT-LLM rc9;
-                # squeeze the batch dim to get (T, H, W, C) for MP4 encoding.
+                # MediaOutput.video is (B, T, H, W, C) since TRT-LLM rc9; accept the
+                # batch-1 5D form or a bare 4D (T, H, W, C). Explicit error (not a bare
+                # assert, which -O strips) on anything unexpected.
                 video = output.video
-                assert (
-                    video.ndim == 5 and video.shape[0] == 1
-                ), f"Expected video shape (1, T, H, W, C), got {video.shape}"
-                frames_np = video[0].cpu().numpy()
+                if video.ndim == 5 and video.shape[0] == 1:
+                    video = video[0]
+                elif video.ndim != 4:
+                    raise RuntimeError(
+                        f"Unexpected video tensor shape {tuple(video.shape)}; "
+                        "expected (T, H, W, C) or (1, T, H, W, C)"
+                    )
+                frames_np = video.cpu().numpy()
                 logger.info(
                     f"Request {request_id}: encoding video output "
                     f"(shape={frames_np.shape}) to MP4 at {fps} fps"
diff --git a/components/src/dynamo/trtllm/workers/video_diffusion_worker.py b/components/src/dynamo/trtllm/workers/video_diffusion_worker.py
index dfca15a416e6..48af300d670e 100644
--- a/components/src/dynamo/trtllm/workers/video_diffusion_worker.py
+++ b/components/src/dynamo/trtllm/workers/video_diffusion_worker.py
@@ -73,6 +73,19 @@ async def init_video_diffusion_worker(
     # Build DiffusionConfig from the main Config
     diffusion_config = DiffusionConfig.from_config(config, skip_components)
 
+    # Make the guidance mode explicit in the logs: dual-guidance (goff) is driven by
+    # config-only flags (nvext can't carry them), so a missing --default-guidance-scale-2
+    # silently falls back to single guidance. Log it prominently instead of hiding it in
+    # the full config dump.
+    logging.info(
+        "Video guidance: %s (g=%s, g2=%s, boundary=%s); cache_backend=%s",
+        "DUAL" if diffusion_config.default_guidance_scale_2 is not None else "SINGLE",
+        diffusion_config.default_guidance_scale,
+        diffusion_config.default_guidance_scale_2,
+        diffusion_config.default_boundary_ratio,
+        diffusion_config.cache_backend,
+    )
+
     # Get the endpoint from the runtime
     endpoint = runtime.endpoint(
         f"{config.namespace}.{config.component}.{config.endpoint}"