perf: TRT zero-copy audit hardening (ControlNet residual, alignment guard, ONNX export)

src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py

@@ -55,10 +55,39 @@ def __init__(self, filepath: str, stream: "cuda.Stream", use_cuda_graph: bool =
         self._fio_cache_len: int = -1  # -1 = not yet initialized
 
         # Sub-phase 5.6: kvo/fio cache inputs are persistent, address-stable,
-        # TRT-contiguous tensors (see models/utils.py create_kvo_cache/create_fi_cache),
-        # so Engine.infer() can bind TensorRT directly to them instead of copying into
-        # its own staging buffer every frame. Rebuilt only when the kvo/fio lazy-init
-        # branches below fire (cache length change), never per-frame.
+        # TRT-contiguous tensors in steady state (models/utils.py
+        # create_kvo_cache/create_fi_cache), so Engine.infer() can bind
+        # TensorRT directly to them instead of copying into its own staging
+        # buffer every frame. `_staging_action` falls back to a copy on
+        # non-contiguous/dtype-mismatch/mis-aligned inputs, and to
+        # bind_and_reset on a pointer change — see utilities.py. Rebuilt only
+        # when the kvo/fio lazy-init branches below fire (cache length
+        # change), never per-frame.
+        #
+        # ControlNet input_control_* residuals are deliberately, PERMANENTLY
+        # EXCLUDED here (reverted from Phase-2 D2 / f8ff50f) and stay on the
+        # DtoD-copy staging path. Two reasons, confirmed against code, not
+        # just bisection:
+        #   1. The residual source toggles between the ControlNet engine's
+        #      live output and idle cached-dummy-zero tensors
+        #      (_add_controlnet_residuals vs _add_cached_dummy_inputs below),
+        #      and a missed bind_and_reset on that toggle silently pins the
+        #      UNet's CUDA graph to whichever buffer was bound at capture
+        #      time.
+        #   2. copy_() runs on the engine stream and orders the ControlNet
+        #      engine's residual write before the UNet graph's read of it; a
+        #      zero-copy bind drops that ordering guarantee with no
+        #      replacement synchronization (see the "copy_() executes on
+        #      engine stream to guarantee ordering" log line in utilities.py).
+        # Re-enabling zero-copy for these names reproduced the "ControlNet
+        # produces no visual change" regression on the rig for both a
+        # single-scale (Canny) and ramped-scale (scribble) config, even
+        # though single-ControlNet residuals are themselves contiguous,
+        # 256-aligned, persistent buffers (controlnet_module.py
+        # build_unet_hook single-CN path) — i.e. the failure is NOT explained
+        # by the copy-fallback path ever firing; it is the ordering guarantee
+        # above. Do not re-add these names to _zero_copy_names without a
+        # rig-verified CUDA-event ordering fix.
         self._zero_copy_names: FrozenSet[str] = frozenset()
 
         self._shape_dict: Dict[str, Any] = {}

src/streamdiffusion/acceleration/tensorrt/utilities.py

@@ -385,8 +385,11 @@ def _apply_gpu_profile_to_config(
     # + EDGE_MASK_CONVOLUTIONS — the sources that produce valid SM_120 kernels.
     # TRT 10.16 exposes TacticSource as an int enum (not IntFlag), so the bitmask
     # is built via (1 << int(source)). No-op on Ada/Ampere.
+    # CUBLAS/CUBLAS_LT are removed from TacticSource in TRT 11 (cuBLAS tactics
+    # dropped entirely) — gate on hasattr so that's a deliberate, logged branch
+    # rather than a silently-swallowed AttributeError.
     if gpu_profile.compute_capability >= (12, 0):
-        try:
+        if hasattr(trt.TacticSource, "CUBLAS"):
             sources = (
                 (1 << int(trt.TacticSource.CUBLAS))
                 | (1 << int(trt.TacticSource.CUBLAS_LT))
@@ -397,8 +400,11 @@ def _apply_gpu_profile_to_config(
             logger.info(
                 "[TRT Config] tactic sources = CUBLAS|CUBLAS_LT|JIT_CONV|EDGE_MASK (CUDNN excluded for SM_120+)"
             )
-        except (AttributeError, TypeError) as e:
-            logger.debug(f"[TRT Config] set_tactic_sources not available: {e}")
+        else:
+            logger.info(
+                "[TRT Config] TRT >=11: CUBLAS/CUBLAS_LT removed from TacticSource — "
+                "default tactic sources already exclude cuDNN/cuBLAS, nothing to scope"
+            )
 
     # max_num_tactics: cap profiling candidates per layer to reduce build time.
     # Available since TRT 10.x; -1 (default) lets TRT decide heuristically. 64 is a
@@ -558,12 +564,28 @@ def _staging_action(
 
     Returns one of:
       "copy"           - copy into self.tensors[name] (today's behavior; always safe)
+      "copy_and_reset" - copy into self.tensors[name], but reset the CUDA graph first
+                         because this name was previously bound zero-copy into a live
+                         graph, which still reads the old bound address rather than
+                         self.tensors[name]
       "bind"           - skip the copy; bind TensorRT directly to the caller's tensor
       "bind_and_reset" - skip the copy and bind directly, but reset the CUDA graph
                          first because the caller's address changed while a graph
                          built from the old address was still live
+
+    A pointer that is not 256-byte aligned falls back to "copy": TensorRT's
+    setTensorAddress contract requires ≥256-byte alignment, and real torch CUDA
+    allocations are always ≥512-byte aligned, so this guard is a defensive
+    invariant against config drift rather than a path exercised in production.
     """
-    if name not in zero_copy_names or not is_contiguous or not dtype_match:
+    if name not in zero_copy_names or not is_contiguous or not dtype_match or cur_ptr % 256 != 0:
+        # Falling back to copy. If this name was bound zero-copy into a LIVE graph,
+        # that graph still reads the stale bound address, not self.tensors[name] —
+        # force one reset so it re-captures reading the staging buffer. prev_ptr is
+        # only non-None for names actually bound before, so plain copy-path inputs
+        # (never in zero_copy_names) never trip this.
+        if graph_exists and prev_ptr is not None:
+            return "copy_and_reset"
         return "copy"
     if graph_exists and cur_ptr != prev_ptr:
         return "bind_and_reset"
@@ -999,12 +1021,8 @@ def get_input_profile_bounds(self, name: str):
         except Exception:
             return None
 
-    def activate(self, reuse_device_memory=None):
-        if reuse_device_memory:
-            self.context = self.engine.create_execution_context_without_device_memory()
-            self.context.device_memory = reuse_device_memory
-        else:
-            self.context = self.engine.create_execution_context()
+    def activate(self):
+        self.context = self.engine.create_execution_context()
 
         # Attach per-layer profiler when STREAMDIFFUSION_PROFILE_TRT is set.
         # Requires engines built with profiling_verbosity=DETAILED for meaningful names.
@@ -1191,8 +1209,14 @@ def infer(self, feed_dict, stream, use_cuda_graph=False, zero_copy_names: frozen
                         cur_ptr,
                         graph_exists,
                     )
-                    if action == "copy":
+                    if action in ("copy", "copy_and_reset"):
                         self.tensors[name].copy_(buf)
+                        if action == "copy_and_reset":
+                            needs_reset = True
+                            # Next frame this name has no prior bind, so
+                            # _staging_action sees prev_ptr=None and returns
+                            # plain "copy" — converges after one reset.
+                            self._bound_ptrs.pop(name, None)
                     else:
                         bind_ptrs[name] = cur_ptr
                         if action == "bind_and_reset":
@@ -1525,7 +1549,6 @@ def export_onnx(
             onnx_path,
             export_params=True,
             opset_version=onnx_opset,
-            do_constant_folding=True,
             input_names=model_data.get_input_names(),
             output_names=model_data.get_output_names(),
             dynamic_axes=model_data.get_dynamic_axes(),
@@ -1577,15 +1600,17 @@ def optimize_onnx(
 ):
     import os
 
-    # Check if external data files exist (indicating external data format was used)
     onnx_dir = os.path.dirname(onnx_path)
-    external_data_files = [f for f in os.listdir(onnx_dir) if f.endswith(".pb")]
-    uses_external_data = len(external_data_files) > 0
+    # Inspect TensorProto.data_location on the raw (unloaded) model rather than
+    # scanning the directory for ".pb" filenames — load_external_data_for_model
+    # resets data_location back to DEFAULT once external data is loaded, so this
+    # check must happen before loading.
+    onnx_model = onnx.load(onnx_path, load_external_data=False)
+    uses_external_data = any(onnx.external_data_helper.uses_external_data(t) for t in onnx_model.graph.initializer)
 
     if uses_external_data:
-        logger.info(f"Optimizing ONNX with external data (found: {external_data_files})")
-        # Load model with external data
-        onnx_model = onnx.load(onnx_path, load_external_data=True)
+        logger.info("Optimizing ONNX with external data")
+        onnx.external_data_helper.load_external_data_for_model(onnx_model, onnx_dir)
         onnx_opt_graph = model_data.optimize(onnx_model)
 
         # Create output directory
@@ -1610,8 +1635,7 @@ def optimize_onnx(
         logger.info("ONNX optimization complete with external data")
 
     else:
-        # Standard optimization for smaller models
-        onnx_model = onnx.load(onnx_path)
+        # No external data to load — the model loaded above is already complete.
         onnx_opt_graph = model_data.optimize(onnx_model)
 
         onnx.save(onnx_opt_graph, onnx_opt_path)

src/streamdiffusion/modules/controlnet_module.py

@@ -13,7 +13,6 @@
 from streamdiffusion.preprocessing.preprocessing_orchestrator import (
     PreprocessingOrchestrator,
 )
-from streamdiffusion.preprocessing.orchestrator_user import OrchestratorUser
 from streamdiffusion.tools.gpu_profiler import profiler
 
 
@@ -89,6 +88,15 @@ def __init__(self, device: str = "cuda", dtype: torch.dtype = torch.float16) ->
         self._cn_cache_images_version: int = -1
         self._cn_cache_scale_hash: Optional[tuple] = None
 
+        # Persistent multi-ControlNet residual merge buffers (Phase-2 prep). The naive
+        # `merged_down[j] = merged_down[j] + ds[j]` allocates a fresh tensor every frame,
+        # which would thrash a CUDA graph if these residuals are later bound zero-copy into
+        # the UNet's input_control_* inputs. Lazily (re)allocated only when residual
+        # shape/dtype/device changes, so pointers stay stable across frames.
+        self._cn_merged_down: Optional[List[torch.Tensor]] = None
+        self._cn_merged_mid: Optional[torch.Tensor] = None
+        self._cn_merged_shape_key: Optional[tuple] = None
+
     # ---------- Public API (used by wrapper in a later step) ----------
     def install(self, stream) -> None:
         self._stream = stream
@@ -116,6 +124,9 @@ def install(self, stream) -> None:
         self._cn_cached_residuals = None
         self._cn_cache_images_version = -1
         self._cn_cache_scale_hash = None
+        self._cn_merged_down = None
+        self._cn_merged_mid = None
+        self._cn_merged_shape_key = None
 
     def add_controlnet(
         self, cfg: ControlNetConfig, control_image: Optional[Union[str, Any, torch.Tensor]] = None
@@ -501,13 +512,6 @@ def _unet_hook(ctx: StepCtx) -> UnetKwargsDelta:
 
             encoder_hidden_states = self._stream.prompt_embeds[:, : self._expected_text_len, :]
 
-            base_kwargs: Dict[str, Any] = {
-                "sample": x_t,
-                "timestep": t_list,
-                "encoder_hidden_states": encoder_hidden_states,
-                "return_dict": False,
-            }
-
             down_samples_list: List[List[torch.Tensor]] = []
             mid_samples_list: List[torch.Tensor] = []
 
@@ -623,16 +627,34 @@ def _unet_hook(ctx: StepCtx) -> UnetKwargsDelta:
                     mid_block_additional_residual=mid_samples_list[0],
                 )
             else:
-                # Merge multiple ControlNet residuals
-                merged_down = down_samples_list[0]
-                merged_mid = mid_samples_list[0]
+                # Merge multiple ControlNet residuals into persistent buffers, in place.
+                # Do NOT alias down_samples_list[0]/mid_samples_list[0] — those are
+                # engine-A's own persistent output buffers. Reallocate only when the
+                # residual shape/dtype/device changes (resolution/batch/model swap),
+                # so the merged tensors stay pointer-stable across frames.
+                shape_key = (
+                    tuple(t.shape for t in down_samples_list[0]),
+                    mid_samples_list[0].shape,
+                    down_samples_list[0][0].dtype,
+                    down_samples_list[0][0].device,
+                )
+                if self._cn_merged_down is None or self._cn_merged_shape_key != shape_key:
+                    self._cn_merged_down = [torch.empty_like(t) for t in down_samples_list[0]]
+                    self._cn_merged_mid = torch.empty_like(mid_samples_list[0])
+                    self._cn_merged_shape_key = shape_key
+
+                for j in range(len(self._cn_merged_down)):
+                    self._cn_merged_down[j].copy_(down_samples_list[0][j])
+                self._cn_merged_mid.copy_(mid_samples_list[0])
+
                 for ds, ms in zip(down_samples_list[1:], mid_samples_list[1:]):
-                    for j in range(len(merged_down)):
-                        merged_down[j] = merged_down[j] + ds[j]
-                    merged_mid = merged_mid + ms
+                    for j in range(len(self._cn_merged_down)):
+                        self._cn_merged_down[j].add_(ds[j])
+                    self._cn_merged_mid.add_(ms)
+
                 _result = UnetKwargsDelta(
-                    down_block_additional_residuals=merged_down,
-                    mid_block_additional_residual=merged_mid,
+                    down_block_additional_residuals=self._cn_merged_down,
+                    mid_block_additional_residual=self._cn_merged_mid,
                 )
 
             # Residual cache write: store result for reuse on upcoming intermediate frames.

src/streamdiffusion/preprocessing/processors/trt_base.py

@@ -132,7 +132,13 @@ def allocate_buffers(self, device: str = "cuda", input_shape: tuple = None):
             if self.engine.get_tensor_mode(name) != trt.TensorIOMode.INPUT:
                 continue
             if input_shape is not None:
-                self.context.set_input_shape(name, input_shape)
+                if not self.context.set_input_shape(name, input_shape):
+                    raise RuntimeError(
+                        f"TensorRTEngine.allocate_buffers: set_input_shape failed for "
+                        f"'{name}' with shape {input_shape}. The engine was built for a "
+                        f"fixed shape range — revert the parameter change or rebuild the "
+                        f"engine for the new shape."
+                    )
             else:
                 static_shape = tuple(self.context.get_tensor_shape(name))
                 if any(d < 0 for d in static_shape):
@@ -141,7 +147,11 @@ def allocate_buffers(self, device: str = "cuda", input_shape: tuple = None):
                         f"shape {static_shape} but no input_shape was provided. "
                         "Pass input_shape=(N, C, H, W) when using a dynamic engine."
                     )
-                self.context.set_input_shape(name, static_shape)
+                if not self.context.set_input_shape(name, static_shape):
+                    raise RuntimeError(
+                        f"TensorRTEngine.allocate_buffers: set_input_shape failed for "
+                        f"'{name}' with shape {static_shape}."
+                    )
 
         # Pass 2: allocate buffers for ALL tensors (output shapes resolved by TRT now).
         for idx in range(self.engine.num_io_tensors):
@@ -192,12 +202,18 @@ def infer(self, feed_dict: dict, stream=None) -> OrderedDict:
                         self.tensors[name] = cached[name]
                 # Re-apply input shapes to TRT context (context state is NOT cached).
                 for name, shape in new_input_shapes.items():
-                    self.context.set_input_shape(name, shape)
+                    if not self.context.set_input_shape(name, shape):
+                        raise RuntimeError(
+                            f"TensorRTEngine.infer: set_input_shape failed for '{name}' with shape {shape}."
+                        )
             else:
                 # LRU miss: reallocate changed inputs, re-derive output shapes.
                 for name, fed_shape in new_input_shapes.items():
                     if fed_shape != tuple(self.tensors[name].shape):
-                        self.context.set_input_shape(name, fed_shape)
+                        if not self.context.set_input_shape(name, fed_shape):
+                            raise RuntimeError(
+                                f"TensorRTEngine.infer: set_input_shape failed for '{name}' with shape {fed_shape}."
+                            )
                         self.tensors[name] = torch.empty(
                             fed_shape,
                             dtype=self.tensors[name].dtype,
@@ -232,7 +248,8 @@ def infer(self, feed_dict: dict, stream=None) -> OrderedDict:
             self.tensors[name].copy_(buf)
 
         for name, tensor in self.tensors.items():
-            self.context.set_tensor_address(name, tensor.data_ptr())
+            if not self.context.set_tensor_address(name, tensor.data_ptr()):
+                raise RuntimeError(f"TensorRTEngine.infer: set_tensor_address failed for '{name}'")
 
         # --- Cross-stream synchronization ---
         # The input copy_() calls above ran on the CURRENT (default) stream.