perf: UNet step sub-phase optimizations (5.0-5.6)#31
Open
forkni wants to merge 1 commit into
Open
Conversation
* refactor: re-apply Phase 1 perf/best-practices cleanup (cosmetic + fail-fast) Re-applies the previously-reverted Phase 1 batch from the perf/best-practices audit: fix the FP8 Q/DQ gate comment to match the actual threshold, convert a dict() call to a dict literal (clears the one pre-existing ruff C408 on a touched file), require an explicit opt_batch_size in EngineBuilder.build() and its 5 compile_* wrappers (real callers already always pass it), drop a dead try/except around a single logger.info call, and tighten two type hints in StreamDiffusion.pipeline (Optional/Union/Tuple). No behavior change except the opt_batch_size fail-fast, which only affects callers that previously relied on an unused default of 1. * refactor: Phase 2 exception hygiene and observability fixes - Add shared _is_oom_error() helper (typed torch.cuda.OutOfMemoryError + string-heuristic fallback) and use it at both TensorRT UNet/VAE engine build fallback sites instead of duplicated inline substring checks. - Log swallowed LoRA candidate-weight-name failures in _load_lora_with_offline_fallback instead of silently continuing. - Log the inner fallback failure in advanced-model-detection instead of a bare except/pass. - Hoist the optional diffusers_ipadapter helper imports in IPAdapterModule.build_unet_hook out of the per-frame hook into a single build-time resolution, and log all previously-silent per-step fallback branches for observability. - Make _load_model fail fast (raise RuntimeError) when an SDXL pipeline retry after a type mismatch also fails, instead of silently continuing with the wrong pipeline type; also null out the stale mismatched `pipe` reference so a later loading-method failure can't let it slip through the final success check. - Add regression tests for _is_oom_error and the SDXL fail-fast path. * chore: migrate deprecated top-level ruff settings to [tool.ruff.lint] Ruff warns that top-level ignore/select/isort/per-file-ignores are deprecated in favour of the lint section. Move them under [tool.ruff.lint] (and .isort / .per-file-ignores); line-length stays at the top level. Verified against cuda-link's pyproject.toml, which already uses this layout. No behavior change: ruff check/format report identical results before and after (same 479 pre-existing repo-wide findings, 0 on the files touched by the Phase 1/2 commits). Investigated adding a matching [tool.pyrefly] section for parity with cuda-link, but did not adopt it: pyrefly falls back to a bundled 'basic' preset when no pyrefly config exists, and that preset is substantially more lenient than explicit config. Adding even an empty [tool.pyrefly] section opts out of it and jumps reported errors from 72 to ~900 (mostly dynamically-set attributes on StreamDiffusion / StreamDiffusionWrapper that were never checked before). Doing this properly needs the same per-module triage cuda-link performed (docs/adr/0005-static-typing-hardening.md) before turning on strict project-includes checking, which is out of scope here. * fix: sync _make_fake_engine test double with TensorRTEngine.__init__ _make_fake_engine (tests/unit/test_trt_engine_guards.py) built engines via TensorRTEngine.__new__() to bypass __init__, but never set _dedicated_stream, _pre_exec_event, _post_exec_event, or _buf_cache -- attributes __init__ has initialized since infer() gained cross-stream sync and LRU shape-cache support. This caused 3 AttributeErrors long treated as "pre-existing, out of scope" failures in every prior phase gate. Not a production bug: __init__ correctly sets all four to None/empty, and infer() guards each with `is not None`. Fix mirrors __init__ in the fake. Full unit suite now fully green (91 passed, 0 failed) with no waived tests. * fix: Phase 3 correctness and resource-safety fixes Four targeted fixes from the perf/best-practices audit, quick-wins #4 and #5: - StreamDiffusion.prepare(): generator defaulted to a pre-constructed torch.Generator(), a mutable default evaluated once at def-time so every instance sharing the default shared one Generator. Default is now None, constructed fresh in the method body. - StreamDiffusionWrapper.prepare(): both runtime stream.prepare() calls (single-prompt and prompt-blending paths) omitted seed, so every runtime prompt change silently reset the RNG to stream.prepare()'s own default. Now forwards seed=getattr(self.stream, "current_seed", 2) to preserve the active seed across prompt changes. - postprocess_image()'s "latent" branch returned an internal decode buffer by reference; callers retaining it across frames would see it mutate. Now clones at this public API boundary. The "np" pinned-buffer alias and the internal __call__/txt2img fast-path aliases are documented as intentional, not cloned (would defeat the pinned-buffer/reuse optimizations). - cleanup_engines_and_rebuild's hardcoded engines_dir = "engines" now honors self._engine_dir when set, matching the existing idiom elsewhere in the file. - cleanup_gpu_memory dropped two explicit Engine.__del__() calls; the method already does del + triple gc.collect() + empty_cache() + ipc_collect() at its tail, and Engine.__del__ is self-guarding, so the explicit dunder calls were redundant. Adds tests/unit/test_phase3_correctness.py (5 regression tests, model-free CPU-only). Full unit suite: 91 passed, 0 failed (no pre-existing failures remain -- see cf3df18). * fix: Phase 4 TensorRT engine-write atomicity and graph/refit robustness Guard TRT engine and timing-cache writes against truncation from an interrupted build via a temp-file + os.replace atomic write helper, matching the idiom already used for FP8 calibration data. Add self-healing recovery around cudaGraphLaunch (reset + fallback to plain execution, re-captures next frame) and a stream-sync guard around the defensive, disabled-by-default Engine.refit() path. Verified: 3 new CPU-only unit tests cover the atomic-write helper's happy path and both failure-preservation cases; a live-GPU smoke run against a cached VAE engine exercised capture, happy-path replay, a forced cudaGraphLaunch failure (recovery + safe fallback, no crash), and self-healed re-capture. * perf: instrument unet_step with profiler.region() spans (Sub-phase 5.0) Adds 5 graph-safe profiler.region() spans (prep/sdxl_cond/hooks/engine/post) inside unet_step to attribute the audit's "~66% unattributed" frame time. Baseline capture (RTX 4090, SDXL-turbo FP8) shows that slice is GPU kernel time inside the single TRT UNet engine call (~93% of frame), not host overhead - every host sub-span measures <0.01ms. Instrumentation only, no behavior change; validated by the existing 94/0 unit suite. * perf: skip redundant set_tensor_address rebind in TRT graph replay (Sub-phase 5.1) Engine.infer() unconditionally re-bound every tensor address via set_tensor_address() before each engine call, even in graphed steady state where the addresses are already baked into the captured CUDA graph (self.tensors buffers are persistent and reused via copy_() - see allocate_buffers/_can_reuse_buffers, which resets cuda_graph_instance to None whenever a buffer is actually reallocated). Skip the rebind loop once a graph is captured; only rebind on first capture or after a reset_cuda_graph() recovery. Verified live on the RTX 4090 td_config.yaml runtime path (the UNet engine loads with use_cuda_graph=True unconditionally via engine_manager.py's loader): 24-frame smoke run with real graph capture + replay shows no errors and frame timing (p50 32.10ms, mean 47.33ms) matching a pre-fix A/B run (p50 31.74ms, mean 47.00ms) - confirms correctness with the expected ~0 measured perf delta (the audit's own 5.0 baseline showed host-side region time is <0.01ms; this is a hygiene/false-dependency fix, not a throughput win). Suite 94/0; ruff clean; pyrefly 72 errors (baseline, no new). * perf: sync-free output/PIL/IPC paths (Sub-phase 5.2) Three output-side sync/allocation fixes, ordered by real-world impact on the TD path: - 5f (only hot-path site): _ipc_pack_rgba / _ipc_pack_unit_rgba now write into a persistent HWC x4 GPU buffer (alpha set once at allocation, B/G/R channels written in place) instead of allocating a fresh torch.full_like alpha + torch.cat every frame. Safe only because the IPC exporters are forced to blocking export (ADR-0001); comments updated at both _lazy_init_ipc_exporter and _lazy_init_cn_ipc_exporter to reflect the buffer is now persistent, not transient. - 5e: _tensor_to_pil_safe (preprocessing_orchestrator.py) reordered CPU-first, following the in-repo base.py:tensor_to_pil template -- the D2H transfer now happens before the tensor.min()/.max() range-check syncs instead of after, collapsing 3 GPU host syncs down to 1. - 5d (hygiene; TD's IPC output path never hits this): _tensor_to_pil_optimized now routes through the shared _output_pin_buf/_d2h_event pinned-buffer + Event machinery instead of a blocking, unpinned .cpu() into pageable memory. Documented the same buffer-reuse caveat as the "np" path (PIL images returned wrap a view of the pinned buffer; callers retaining them across frames must .copy()). Verified via tests/unit/test_sync_free_output_5_2.py: byte-identical parity against frozen pre-5.2 reference formulas for all three sites, plus persistent-buffer reuse/realloc and independent-buffer-identity checks (10 new tests, GPU-gated). Suite 104/0 (94 baseline + 10 new); ruff clean on touched files; pyrefly 72 errors (baseline, no new). * fix: batch-dynamic engines skip l2tc tiling, silencing benign VALIDATE FAIL build_engine() passed dynamic_shapes=build_dynamic_shape into Engine.build(), which tracks only dynamic resolution. On the default "Flexible" preset the UNet is resolution-static but batch-dynamic (build_static_batch=False), so dynamic_shapes was passed False, and _apply_gpu_profile_to_config's tiling branch ran against a graph that still has a symbolic batch dim -- TRT emitted "[l2tc] VALIDATE FAIL - Graph contains symbolic shape" as a no-op for every applicable layer, for the entire ~11.5 min production TD session analyzed this round. Fix: dynamic_shapes=build_dynamic_shape or not build_static_batch (one line, utilities.py:1372) -- True whenever any dim, resolution or batch, is symbolic, matching l2tc's actual requirement (ALL dims must be concrete). Fully-static engines (e.g. ControlNet: build_static_batch=True, build_dynamic_shape=False) are unaffected -- still get tiling as before. Tightened the _apply_gpu_profile_to_config docstring to match. Verified via new tests/unit/test_l2tc_dynamic_shapes.py: calls build_engine() with GPU detection/memory query/the real TRT build all monkeypatched out (no CUDA/TRT context required) and asserts the dynamic_shapes kwarg it wires into Engine.build() for the previously-broken case (batch-dynamic, resolution-static -> True), the unaffected fully-static case (-> False), and the already-working resolution-dynamic case (-> True). Suite 107/0 (104 baseline + 3 new); ruff clean; pyrefly 72 errors (baseline, no new). Impact is log-spam-only -- zero perf/numerical change, visible on the next engine rebuild. * perf: elide trt.input_staging DtoD copy for kvo/fio UNet cache inputs (Sub-phase 5.6) Engine.infer() copied every feed_dict tensor into its own persistent staging buffer each frame solely to give TensorRT a stable contiguous address to bind. nsys profiling (5.0/5.5) attributed ~9% of frame time (p95 3.1ms, ~161 DtoD copies/frame) to this region on the UNet call, dominated by the kvo/fio StreamV2V cache inputs. Those cache tensors are already persistent, address-stable, TRT-contiguous buffers (create_kvo_cache/create_fi_cache build them with layers_in_bucket as the outermost dim specifically for this) — the only event that moves their address is a batch-size change, which already forces a full buffer realloc+graph reset today. So copying them every frame is pure waste. Add an opt-in zero_copy_names parameter to Engine.infer() and a pure _staging_action() decision helper (copy/bind/bind_and_reset) that lets TensorRT bind directly to the caller's tensor instead of copying into self.tensors[name], guarded by contiguity and dtype-match fallbacks. Wire it in UNet2DConditionModelEngine for the kvo/fio input names only; every other caller (VAE, ControlNet, safety, preprocessing) passes nothing and keeps the original copy path byte-for-byte. * fix: lower ControlNet dynamic-shape floor 384->256 via shared min_image_shape (bug #6) ControlNet's dynamic-shape get_input_profile() hardcoded a 384px floor independent of the UNet's 256px floor (BaseModel.min_image_shape), so any resolution in [256, 384) hard-failed at Engine.allocate_buffers the moment ControlNet was enabled. Derive the ControlNet floor/ceiling from the same BaseModel.min_image_shape/max_image_shape the UNet already uses (SDXL inherits via super()), lower the matching engine_manager.py build-option default, and relabel the ControlNet dynamic-engine cache directory (--dyn-384-1024 -> --dyn-256-1024) so a re-floored build never collides with a stale 384-floored cache. Trade-off: widening the dynamic profile can slightly reduce TRT tactic specialization at the optimum resolution -- acceptable given the alternative is a hard failure. Co-Authored-By: Claude Sonnet 5 <[email protected]> * perf: normalize blend weights on CPU + persistent ControlNet conditioning_scale scalar (Sub-phase 5.4) Two per-frame GPU creation syncs removed: - _normalize_weights built a CUDA tensor from a Python list every call, but every consumer (.item()/.tolist() in slerp/multi_slerp/cosine_weighted paths, float*cuda_tensor in the linear paths) immediately reads it back or treats it as a scalar operand. Build it on CPU float32 instead -- same consumers work unchanged (a 0-dim CPU tensor is treated as a wrapped scalar when multiplied against a CUDA tensor), the sync and readback are both gone, and float32 is more precise than the model's fp16 for the normalization divide. - ControlNetModelEngine rebuilt a conditioning_scale tensor via torch.tensor(...) on every __call__. Replace with one lazily-allocated persistent scalar updated via fill_() (an async kernel launch, no sync), mirroring the existing _fi_strength_tensor.fill_() idiom. Each active ControlNet has its own engine instance, so a per-instance buffer is correct; the binding rank (_cs_rank1) is fixed at load time so the shape never changes underneath it. Also drops three dead variable assignments and fixes one out-of-order import in stream_parameter_updater.py (pre-existing lint debt in a file this change already touches; the repo's commit gate lints the whole staged file, not just the diff). Co-Authored-By: Claude Sonnet 5 <[email protected]> * perf: 1-frame-delayed async NSFW safety-checker readback (Sub-phase 5.3) Removes the cudaStreamSynchronize that NSFWDetectorEngine.__call__ forced on the hot path via probs[0,0].item() whenever use_safety_checker=True. The classifier now writes its GPU probability into a pinned host scalar via a non_blocking copy; _apply_safety_checker reads the *previous* frame's pinned verdict before launching this frame's classification, mirroring the existing SimilarImageFilter 1-frame-delayed readback pattern (image_filter.py). Edge cases: first frame passes through unconditionally (no prior verdict to act on); a 1-frame detection-leak window is possible before flagging lands. Both are acceptable for an opt-in, off-by-default, TensorRT-only feature. Updated tests/unit/test_safety_checker.py for the new (tensor, prob_pin) -> None checker contract and the delayed-readback semantics (10/10 passing). Co-Authored-By: Claude Sonnet 5 <[email protected]> * chore: add Claude review workflow to PR2 branch for head-ref validation * chore: widen Claude review allowedTools (Skill, python/pytest, /tmp Write) to unblock permission-denied max-turns failures * chore: raise Claude review max-turns to 60, discourage filesystem-wide find and output redirection * fix: correct NSFW verdict-to-frame attribution and guard pin_memory on CPU-only builds _apply_safety_checker previously read the prior frame's async verdict but applied it to the current frame's pixels, letting an isolated NSFW frame leak through and wrongly blanking the next clean frame. Buffer the raw frame in _pending_frame and emit it once its own verdict lands (+1 frame latency, first call now emits a black startup frame instead of raw pixels). Also guard .pin_memory() behind torch.cuda.is_available() so the CPU-only test contract in test_safety_checker.py holds. --------- Co-authored-by: Claude Sonnet 5 <[email protected]>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
wrapper.py.test_safety_checker.pyupdates.forkni/StreamDiffusion@SDTD_040_beta_release(internal PR feat: port kvo_cache patch onto diffusers 0.38.0 (runtime monkey-patch) #13).Test plan
SDTD_040_beta_releasetip, no conflicts.src/streamdiffusion/wrapper.py,tests/unit/test_safety_checker.py).