diff --git a/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py b/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py index de317065d..85d081c7a 100644 --- a/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py +++ b/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] = {} diff --git a/src/streamdiffusion/acceleration/tensorrt/utilities.py b/src/streamdiffusion/acceleration/tensorrt/utilities.py index 7cea0b456..caea2daa2 100644 --- a/src/streamdiffusion/acceleration/tensorrt/utilities.py +++ b/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) diff --git a/src/streamdiffusion/modules/controlnet_module.py b/src/streamdiffusion/modules/controlnet_module.py index 2fc6ef520..03c57345d 100644 --- a/src/streamdiffusion/modules/controlnet_module.py +++ b/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. diff --git a/src/streamdiffusion/preprocessing/processors/trt_base.py b/src/streamdiffusion/preprocessing/processors/trt_base.py index 56b8852a6..b97778b21 100644 --- a/src/streamdiffusion/preprocessing/processors/trt_base.py +++ b/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. diff --git a/src/streamdiffusion/tools/compile_depth_anything_tensorrt.py b/src/streamdiffusion/tools/compile_depth_anything_tensorrt.py index ac0ed7465..b842ec101 100644 --- a/src/streamdiffusion/tools/compile_depth_anything_tensorrt.py +++ b/src/streamdiffusion/tools/compile_depth_anything_tensorrt.py @@ -9,14 +9,15 @@ python -m streamdiffusion.tools.compile_depth_anything_tensorrt --model_size small --resolution 518 """ -import torch +import importlib.util import logging -import os from pathlib import Path -from typing import Optional + import fire +import torch -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') + +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) try: @@ -24,18 +25,13 @@ from streamdiffusion.acceleration.tensorrt.utilities import BUILD_TRT_LOGGER - TENSORRT_AVAILABLE = True except ImportError: TENSORRT_AVAILABLE = False BUILD_TRT_LOGGER = None logger.warning("TensorRT not available. Please install it first.") -try: - import onnx - ONNX_AVAILABLE = True -except ImportError: - ONNX_AVAILABLE = False +ONNX_AVAILABLE = importlib.util.find_spec("onnx") is not None # Depth Anything model configs DEPTH_ANYTHING_MODELS = { @@ -55,10 +51,7 @@ def export_depth_anything_to_onnx( - onnx_path: Path, - model_size: str = "small", - resolution: int = 518, - device: str = "cuda" + onnx_path: Path, model_size: str = "small", resolution: int = 518, device: str = "cuda" ) -> bool: """Export Depth Anything model to ONNX format""" try: @@ -107,6 +100,7 @@ def export_depth_anything_to_onnx( except Exception as e: logger.error(f"Failed to export ONNX: {e}") import traceback + traceback.print_exc() return False @@ -126,11 +120,11 @@ def build_tensorrt_engine( try: builder = trt.Builder(BUILD_TRT_LOGGER) - network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + network = builder.create_network() # EXPLICIT_BATCH deprecated/ignored in TRT 10.x parser = trt.OnnxParser(network, BUILD_TRT_LOGGER) # Parse ONNX - with open(onnx_path, 'rb') as f: + with open(onnx_path, "rb") as f: if not parser.parse(f.read()): for i in range(parser.num_errors): logger.error(f"ONNX parse error: {parser.get_error(i)}") @@ -146,10 +140,12 @@ def build_tensorrt_engine( # Set optimization profile for fixed resolution profile = builder.create_optimization_profile() - profile.set_shape("input", - (1, 3, resolution, resolution), # min - (1, 3, resolution, resolution), # opt - (1, 3, resolution, resolution)) # max + profile.set_shape( + "input", + (1, 3, resolution, resolution), # min + (1, 3, resolution, resolution), # opt + (1, 3, resolution, resolution), + ) # max config.add_optimization_profile(profile) # Build engine @@ -162,7 +158,7 @@ def build_tensorrt_engine( # Save engine engine_path.parent.mkdir(parents=True, exist_ok=True) - with open(engine_path, 'wb') as f: + with open(engine_path, "wb") as f: f.write(serialized_engine) logger.info(f"TensorRT engine saved: {engine_path}") @@ -171,6 +167,7 @@ def build_tensorrt_engine( except Exception as e: logger.error(f"Failed to build TensorRT engine: {e}") import traceback + traceback.print_exc() return False @@ -210,7 +207,7 @@ def compile_depth_anything( # Check if engine already exists if engine_path.exists(): logger.info(f"Engine already exists: {engine_path}") - overwrite = input("Overwrite? (y/N): ").lower().strip() == 'y' + overwrite = input("Overwrite? (y/N): ").lower().strip() == "y" if not overwrite: return @@ -232,9 +229,9 @@ def compile_depth_anything( logger.info("Removed intermediate ONNX file") logger.info(f"\nSuccess! Engine saved to: {engine_path}") - logger.info(f"\nTo use in config:") - logger.info(f' preprocessor: "depth_tensorrt"') - logger.info(f' preprocessor_params:') + logger.info("\nTo use in config:") + logger.info(' preprocessor: "depth_tensorrt"') + logger.info(" preprocessor_params:") logger.info(f' engine_path: "{engine_path}"') diff --git a/tests/unit/test_controlnet_residual_merge.py b/tests/unit/test_controlnet_residual_merge.py new file mode 100644 index 000000000..d9b1224e1 --- /dev/null +++ b/tests/unit/test_controlnet_residual_merge.py @@ -0,0 +1,186 @@ +""" +Regression tests for ControlNetModule's multi-ControlNet residual merge (Phase-2 prep, D1). + +CPU-only and model-free: fake ControlNet callables stand in for real engines so the +merge logic in build_unet_hook()'s closure can be exercised without CUDA/TRT. + +Root cause being guarded: the old merge (`merged_down[j] = merged_down[j] + ds[j]`) +allocated a fresh tensor every frame, so the UNet's input_control_* residuals would +never be pointer-stable across frames — a prerequisite for zero-copy binding those +inputs (Phase-2 D2). This also verifies the merge never aliases engine A's own +persistent output buffer. +""" + +import torch + +from streamdiffusion.hooks import StepCtx +from streamdiffusion.modules.controlnet_module import ControlNetModule + + +class _FakeCN: + """Stands in for a ControlNet engine: returns fixed-value residuals, records calls.""" + + def __init__(self, down_value: float, mid_value: float, down_shapes, mid_shape): + self.calls = 0 + self._down_value = down_value + self._mid_value = mid_value + self._down_shapes = down_shapes + self._mid_shape = mid_shape + + def __call__( + self, sample, timestep, encoder_hidden_states, controlnet_cond, conditioning_scale, return_dict=False + ): + self.calls += 1 + down = [torch.full(shape, self._down_value, dtype=torch.float32) for shape in self._down_shapes] + mid = torch.full(self._mid_shape, self._mid_value, dtype=torch.float32) + return down, mid + + +class _FakeStream: + def __init__(self, text_len: int = 77, batch: int = 1): + self.prompt_embeds = torch.randn(batch, text_len, 8) + + +def _make_module_with_two_controlnets(down_shapes, mid_shape) -> tuple: + module = ControlNetModule(device="cpu", dtype=torch.float32) + module._stream = _FakeStream() + + cn_a = _FakeCN(down_value=1.0, mid_value=2.0, down_shapes=down_shapes, mid_shape=mid_shape) + cn_b = _FakeCN(down_value=3.0, mid_value=4.0, down_shapes=down_shapes, mid_shape=mid_shape) + + module.controlnets = [cn_a, cn_b] + module.controlnet_images = [torch.randn(1, 3, 8, 8), torch.randn(1, 3, 8, 8)] + module.controlnet_scales = [1.0, 1.0] + module.enabled_list = [True, True] + + return module, cn_a, cn_b + + +def _make_ctx(batch: int = 1) -> StepCtx: + return StepCtx( + x_t_latent=torch.randn(batch, 4, 8, 8), + t_list=torch.tensor([0]), + step_index=0, + guidance_mode="none", + sdxl_cond=None, + ) + + +class TestControlNetResidualMerge: + def test_merge_is_numerically_correct(self): + down_shapes = [(1, 4, 8, 8), (1, 4, 4, 4)] + mid_shape = (1, 4, 2, 2) + module, cn_a, cn_b = _make_module_with_two_controlnets(down_shapes, mid_shape) + hook = module.build_unet_hook() + + result = hook(_make_ctx()) + + assert cn_a.calls == 1 and cn_b.calls == 1 + for merged, shape in zip(result.down_block_additional_residuals, down_shapes): + assert torch.allclose(merged, torch.full(shape, 4.0)), "1.0 + 3.0 == 4.0 per down block" + assert torch.allclose(result.mid_block_additional_residual, torch.full(mid_shape, 6.0)), "2.0 + 4.0 == 6.0" + + def test_merge_buffers_are_pointer_stable_across_frames(self): + """Same shape on consecutive frames must reuse the same tensor objects + (required before the UNet's input_control_* inputs can be zero-copy bound).""" + down_shapes = [(1, 4, 8, 8)] + mid_shape = (1, 4, 2, 2) + module, _cn_a, _cn_b = _make_module_with_two_controlnets(down_shapes, mid_shape) + hook = module.build_unet_hook() + + result1 = hook(_make_ctx()) + down_ptr1 = [t.data_ptr() for t in result1.down_block_additional_residuals] + mid_ptr1 = result1.mid_block_additional_residual.data_ptr() + + result2 = hook(_make_ctx()) + down_ptr2 = [t.data_ptr() for t in result2.down_block_additional_residuals] + mid_ptr2 = result2.mid_block_additional_residual.data_ptr() + + assert down_ptr1 == down_ptr2, "merged down-block buffers must be reused, not reallocated" + assert mid_ptr1 == mid_ptr2, "merged mid-block buffer must be reused, not reallocated" + # Values still correct on frame 2 (buffer was overwritten, not just left stale) + assert torch.allclose(result2.down_block_additional_residuals[0], torch.full(down_shapes[0], 4.0)) + + def test_merge_reallocates_on_shape_change(self): + """A resolution/batch change must produce a new buffer, not corrupt-reuse the old one.""" + down_shapes = [(1, 4, 8, 8)] + mid_shape = (1, 4, 2, 2) + module, cn_a, cn_b = _make_module_with_two_controlnets(down_shapes, mid_shape) + hook = module.build_unet_hook() + + result1 = hook(_make_ctx()) + down_ptr1 = result1.down_block_additional_residuals[0].data_ptr() + + # Simulate a resolution change: new shapes for both fake engines. + new_down_shapes = [(1, 4, 16, 16)] + new_mid_shape = (1, 4, 4, 4) + cn_a._down_shapes = new_down_shapes + cn_a._mid_shape = new_mid_shape + cn_b._down_shapes = new_down_shapes + cn_b._mid_shape = new_mid_shape + + result2 = hook(_make_ctx()) + assert result2.down_block_additional_residuals[0].shape == new_down_shapes[0] + down_ptr2 = result2.down_block_additional_residuals[0].data_ptr() + assert down_ptr2 != down_ptr1, "shape change must trigger reallocation" + + def test_merge_does_not_alias_engine_output_buffer(self): + """The merge must not mutate/alias down_samples_list[0] — that's engine A's + own persistent output buffer, reused by the engine on the next frame.""" + down_shapes = [(1, 4, 8, 8)] + mid_shape = (1, 4, 2, 2) + module, cn_a, _cn_b = _make_module_with_two_controlnets(down_shapes, mid_shape) + hook = module.build_unet_hook() + + result = hook(_make_ctx()) + + # Re-invoke cn_a directly (as the engine would on its own next call) and confirm + # its freshly-returned buffer is untouched by the merge (still all 1.0, not 4.0). + fresh_down, fresh_mid = cn_a( + sample=None, timestep=None, encoder_hidden_states=None, controlnet_cond=None, conditioning_scale=1.0 + ) + assert torch.allclose(fresh_down[0], torch.full(down_shapes[0], 1.0)) + assert torch.allclose(fresh_mid, torch.full(mid_shape, 2.0)) + # Sanity: the merge result itself is still the summed value. + assert torch.allclose(result.down_block_additional_residuals[0], torch.full(down_shapes[0], 4.0)) + + def test_single_controlnet_bypasses_merge_buffers(self): + """With only one active ControlNet, the engine's own output is returned + directly — no merge buffer should be allocated.""" + module = ControlNetModule(device="cpu", dtype=torch.float32) + module._stream = _FakeStream() + down_shapes = [(1, 4, 8, 8)] + mid_shape = (1, 4, 2, 2) + cn_a = _FakeCN(down_value=5.0, mid_value=6.0, down_shapes=down_shapes, mid_shape=mid_shape) + module.controlnets = [cn_a] + module.controlnet_images = [torch.randn(1, 3, 8, 8)] + module.controlnet_scales = [1.0] + module.enabled_list = [True] + + hook = module.build_unet_hook() + hook(_make_ctx()) + + assert module._cn_merged_down is None + assert module._cn_merged_mid is None + + def test_install_resets_merge_buffers(self): + down_shapes = [(1, 4, 8, 8)] + mid_shape = (1, 4, 2, 2) + module, _cn_a, _cn_b = _make_module_with_two_controlnets(down_shapes, mid_shape) + module.build_unet_hook()(_make_ctx()) + assert module._cn_merged_down is not None + + class _MinimalStream: + unet_hooks = [] + controlnets = None + controlnet_scales = None + preprocessors = None + + # attach_orchestrator requires a preprocessing orchestrator; reuse the existing one + # (install() only touches it when _preprocessing_orchestrator is None). + module._preprocessing_orchestrator = object() + module.install(_MinimalStream()) + + assert module._cn_merged_down is None + assert module._cn_merged_mid is None + assert module._cn_merged_shape_key is None diff --git a/tests/unit/test_zero_copy_staging_5_6.py b/tests/unit/test_zero_copy_staging_5_6.py index 56bbe341f..dbba00bbf 100644 --- a/tests/unit/test_zero_copy_staging_5_6.py +++ b/tests/unit/test_zero_copy_staging_5_6.py @@ -34,8 +34,9 @@ ) -PTR_A = 0xDEAD_BEEF -PTR_B = 0xFEED_FACE +PTR_A = 0x1000_0000 # 256-byte aligned (audit M2 guard) +PTR_B = 0x2000_0000 # 256-byte aligned (audit M2 guard) +PTR_MISALIGNED = PTR_A + 1 # not 256-byte aligned class TestStagingActionNotZeroCopy: @@ -54,13 +55,18 @@ def test_name_not_in_zero_copy_names_copies(self): assert result == "copy" def test_empty_zero_copy_names_always_copies(self): - """Default frozenset() → today's behavior exactly, for every name.""" + """Default frozenset() → today's behavior exactly, for every name. + + prev_ptr=None is the realistic value here: a name that is never in + zero_copy_names is never bound, so Engine._bound_ptrs never has an + entry for it and _bound_ptrs.get(name) is always None. + """ result = _staging_action( "kvo_cache_in_0", frozenset(), True, True, - PTR_A, + None, PTR_A, True, ) @@ -179,3 +185,149 @@ def test_first_bind_with_no_prior_ptr_recorded(self): False, ) assert result == "bind" + + +class TestStagingActionAlignmentGuard: + """Audit M2: setTensorAddress requires >=256-byte alignment. A mis-aligned + cur_ptr must fall back to "copy" even when every other eligibility check + would otherwise allow a bind — never silently bind an unaligned address.""" + + def test_misaligned_pointer_falls_back_to_copy_no_graph(self): + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + True, + True, + None, + PTR_MISALIGNED, + False, + ) + assert result == "copy" + + def test_misaligned_pointer_falls_back_to_copy_with_live_graph(self): + """A name previously bound zero-copy (prev_ptr set from an earlier, + valid aligned bind) whose current address is invalid for + setTensorAddress must copy — but a live graph is still replaying the + stale bound address, so this must ALSO force a reset. + + This is the regression lock: f8ff50f's bind->copy fallback left the + live UNet CUDA graph replaying a frozen buffer forever once a + zero-copy input fell back to copy mid-stream (see unet_engine.py's + header comment for the full incident writeup). Before the + copy_and_reset outcome existed, this exact transition asserted plain + "copy" here — i.e. the test documented the bug as correct behavior. + """ + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + True, + True, + PTR_A, # previously bound at a valid, aligned address + PTR_MISALIGNED, # now invalid -> must fall back to copy + True, + ) + assert result == "copy_and_reset" + + +class TestStagingActionCopyAndReset: + """copy_and_reset is the general-purpose safety net for ANY zero-copy name + (currently kvo/fio only — see unet_engine.py) that falls back to copy while + a graph built from its previously-bound address is still live. Exercises + the non-contiguous fallback path specifically (the misaligned-pointer path + is covered by TestStagingActionAlignmentGuard above).""" + + def test_non_contiguous_previously_bound_with_live_graph_copies_and_resets(self): + """A previously-bound name (e.g. a cache tensor that got sliced/viewed + this frame) that is no longer contiguous, with a live graph, must copy + AND force a reset — the graph is still reading the old bound address.""" + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + False, # is_contiguous + True, + PTR_A, # previously bound + PTR_A, + True, + ) + assert result == "copy_and_reset" + + def test_non_contiguous_previously_bound_no_live_graph_just_copies(self): + """Same fallback, but no graph exists yet — nothing to invalidate, so a + plain copy is correct; forcing a reset here would be a needless no-op + reset on every non-graphed frame.""" + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + False, + True, + PTR_A, + PTR_A, + False, + ) + assert result == "copy" + + def test_never_bound_name_falls_back_with_live_graph_just_copies(self): + """prev_ptr=None means this name has never been bound zero-copy, so no + graph could possibly be reading a stale address for it — plain copy, + no spurious reset.""" + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + False, + True, + None, + PTR_A, + True, + ) + assert result == "copy" + + +class TestStagingActionControlNetResiduals: + """Phase-2 D2 attempted to opt input_control_* UNet inputs into this same + zero-copy path; that was reverted (see unet_engine.py header comment) after + it reproduced a "ControlNet produces no visual change" regression on the + rig, driven by a lost engine-stream copy ordering guarantee — not by the + bind->copy fallback covered above. input_control_* names are PERMANENTLY + excluded from the real UNet engine's _zero_copy_names. The tests below + exercise _staging_action as a pure function using ControlNet-shaped names + only as convenient example zero-copy names; they do not describe current + runtime behavior for actual ControlNet residuals.""" + + def test_steady_state_control_residual_binds_without_reset(self): + """Same persistent merge/output buffer across frames — plain bind.""" + result = _staging_action( + "input_control_00", + frozenset({"input_control_00"}), + True, + True, + PTR_A, + PTR_A, + True, + ) + assert result == "bind" + + def test_idle_active_toggle_binds_and_resets(self): + """Idle dummy-zero buffer -> active ControlNet residual buffer address + flip while a graph is live must force a re-capture.""" + result = _staging_action( + "input_control_00", + frozenset({"input_control_00"}), + True, + True, + PTR_A, + PTR_B, + True, + ) + assert result == "bind_and_reset" + + def test_middle_control_residual_binds_without_reset(self): + result = _staging_action( + "input_control_middle", + frozenset({"input_control_middle"}), + True, + True, + PTR_A, + PTR_A, + True, + ) + assert result == "bind"