From 7e1bc5cbde862f33bf07ec81e244704cccebd451 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 17:33:01 +0000 Subject: [PATCH] Tighten fault-path comments added by the video/diffusion hardening pass --- scripts/fp8_layer_ablation.py | 5 ++--- scripts/sdpa_mask_backend_probe.py | 3 +-- scripts/video_speedmem_bench.py | 4 ++-- .../backend/core/inference/diffusion_attention.py | 2 +- studio/backend/core/inference/diffusion_cache.py | 14 ++++++-------- .../core/inference/diffusion_cfg_parallel.py | 12 ++++-------- studio/backend/core/inference/video.py | 15 ++++++--------- 7 files changed, 22 insertions(+), 33 deletions(-) diff --git a/scripts/fp8_layer_ablation.py b/scripts/fp8_layer_ablation.py index 71e52ebd60..1cd8e1e471 100644 --- a/scripts/fp8_layer_ablation.py +++ b/scripts/fp8_layer_ablation.py @@ -452,9 +452,8 @@ def main(argv = None) -> int: ) t0 = time.perf_counter() pipe = _build_pipe(repo, force_fp32) - # _build_pipe returns a CPU pipeline (the bench applies levers first, then places on CUDA). - # The ablation captures a forward directly with a CUDA generator and reads the DiT off the - # GPU below, so place the pipeline on CUDA here rather than crashing on a cpu/cuda mismatch. + # _build_pipe returns a CPU pipeline; the ablation captures a forward with a CUDA generator, + # so place it on CUDA here to avoid a cpu/cuda mismatch. pipe = pipe.to("cuda") print(f"[load] pipe built in {time.perf_counter()-t0:.1f}s", flush = True) diff --git a/scripts/sdpa_mask_backend_probe.py b/scripts/sdpa_mask_backend_probe.py index f9d3f71c63..0dfd25ea7c 100644 --- a/scripts/sdpa_mask_backend_probe.py +++ b/scripts/sdpa_mask_backend_probe.py @@ -30,8 +30,7 @@ def timed(fn, iters = 20): torch.cuda.synchronize() return (time.perf_counter() - t0) / iters * 1e3 except torch.OutOfMemoryError: - # An occupied / too-small cuda:0 OOMs on the dense NxN mask; that is a memory limit, - # not a backend rejecting the mask, so don't mislabel it UNSUPPORTED. + # OOM on the dense NxN mask is a memory limit, not a backend rejecting it; don't mislabel UNSUPPORTED. torch.cuda.empty_cache() return "OOM" except Exception as e: # noqa: BLE001 diff --git a/scripts/video_speedmem_bench.py b/scripts/video_speedmem_bench.py index 7aa10b5006..f15ff9d6e5 100644 --- a/scripts/video_speedmem_bench.py +++ b/scripts/video_speedmem_bench.py @@ -623,8 +623,8 @@ def _timed_video( marker = getattr(transformer, "_unsloth_step_cache", None) if not marker or str(marker).endswith(f"#s{int(steps)}"): return "magcache" # already sized for these steps - # Fail closed like production: reapplying over a cache that would not disengage - # double-hooks the transformer and times a stale/stacked curve as if it were fresh. + # Fail closed: reapplying over a cache that would not disengage double-hooks the + # transformer and times a stale curve as if it were fresh. if not _disengage_step_cache( transformer, reason = f"explicit magcache re-interpolating for {steps} steps", diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index feb9a26060..7ce62c8f3c 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -64,7 +64,7 @@ def normalize_attention_backend(value: Optional[str]) -> Optional[str]: # mid-generation). Gate by a (min, max-exclusive) capability range: FA3 is Hopper-SM90 only # (upper bound, so flash3 on a B200 drops to native), FA4 is Blackwell+ (no upper bound). _ARCH_CAPABILITY: dict[str, tuple[tuple[int, int], Optional[tuple[int, int]]]] = { - "flash": ((8, 0), None), # Dao-AILab FlashAttention 2 -> Ampere (SM80)+ (no Turing) + "flash": ((8, 0), None), # FlashAttention 2 -> Ampere (SM80)+ "_flash_3_hub": ((9, 0), (10, 0)), # FlashAttention 3 -> Hopper (SM90) only "flash_4_hub": ((10, 0), None), # FlashAttention 4 -> Blackwell (SM100)+ } diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index 7375624274..dfc3570329 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -681,9 +681,8 @@ def apply_step_cache( disable_cache() transformer._unsloth_step_cache = None except Exception as rollback_exc: # noqa: BLE001 - # Both the enable AND its cleanup failed: the transformer may keep partial hooks - # while we'd otherwise report a clean uncached None. Surface it so the caller - # reloads instead of generating on a half-cached model. + # Enable and its cleanup both failed: surface it so the caller reloads instead + # of generating on a half-cached model. raise RuntimeError( "step-cache enable failed and rollback also failed; the transformer may be " "partially cached and must be reloaded " @@ -781,9 +780,8 @@ def maybe_toggle_step_cache( # endswith, not substring: "#s5" would match inside "#s50". and not str(engaged).endswith(f"#s{int(steps)}") ): - # A failed removal used to short-circuit and fall through to `return mode`, reporting - # "magcache" while the OLD #sN curve stayed armed (wrong ratio schedule, silently - # degraded output). Fail closed so the caller reloads instead. + # Fail closed: a failed removal would leave the old #sN curve armed while reporting + # "magcache" with the wrong ratio schedule. if not _disengage_step_cache( transformer, reason = f"magcache re-interpolating for {steps} steps", logger = logger ): @@ -805,8 +803,8 @@ def maybe_toggle_step_cache( logger = logger, ) if not want and engaged: - # Below the cache threshold we want uncached; a failed disable leaves the (possibly - # wrong-step) cache armed, so surface it rather than reporting the stale mode. + # Below the threshold we want uncached; a failed disable leaves the cache armed, so + # surface it instead of the stale mode. if not _disengage_step_cache( transformer, reason = f"auto: {steps} steps < {FBCACHE_MIN_STEPS}", diff --git a/studio/backend/core/inference/diffusion_cfg_parallel.py b/studio/backend/core/inference/diffusion_cfg_parallel.py index 58485409e7..280a4042cc 100644 --- a/studio/backend/core/inference/diffusion_cfg_parallel.py +++ b/studio/backend/core/inference/diffusion_cfg_parallel.py @@ -230,10 +230,8 @@ class CFGParallelProxy: raise def disable_cache(self) -> None: - # Removal must be transactional: if the primary's disable_cache raised while it ran - # outside this guard, the replica was never cleaned and _broken stayed False, so a - # half-removed pair kept routing. Attempt BOTH, record every failure, and only then - # decide _broken -- any failure disables routing and surfaces so the caller reloads. + # Transactional removal: attempt BOTH branches, record every failure, then decide + # _broken. A primary-only failure previously left the replica cached and kept routing. failures: list[tuple[str, Exception]] = [] for name, module in (("primary", self._primary), ("replica", self._replica)): try: @@ -412,10 +410,8 @@ def _install_threadsafe_cudnn_attention(logger: Any = None) -> bool: return_lse = return_lse, _parallel_config = _parallel_config, ) - # F.scaled_dot_product_attention (what the stock backend calls) treats a boolean - # mask as "True participates" and converts it to an ADDITIVE bias internally; the - # lower-level cuDNN op takes that bias directly, so a bool mask passed straight - # through is misread for any partial (non-all-True) mask. Convert to match SDPA. + # SDPA converts a bool mask ("True participates") to an additive bias internally; + # the cuDNN op takes the bias directly, so convert to match it for partial masks. if attn_mask is not None and attn_mask.dtype == torch.bool: attn_mask = torch.zeros_like(attn_mask, dtype = query.dtype).masked_fill_( ~attn_mask, float("-inf") diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 99350a052d..1a196504dd 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -457,10 +457,9 @@ def _step_cache_all_or_none( for view, expert_name in pairs: results.append((view, expert_name, engage_fn(view, expert_name))) except BaseException as exc: - # A later expert raising mid-loop leaves the experts engaged BEFORE it still cached - # while the load unwinds -- the same silent all-or-none violation as a mixed outcome, - # so tear down every expert that got a marker, then re-raise (or a reload-required - # error if rollback itself fails). + # A later expert raising mid-loop leaves earlier experts cached (an all-or-none + # violation), so tear down every marked expert then re-raise, or raise reload-required + # if rollback itself fails. rollback_failed: list[str] = [] for view, name in pairs: transformer = getattr(view, "transformer", None) @@ -645,9 +644,8 @@ class VideoBackend: normalize_te_quant(text_encoder_quant) # Same for vae_quant (the dense VAE is resident for every load kind). normalize_vae_quant(vae_quant) - # Reject malformed cache-quality / cfg-parallel here too: the HTTP Literal fields gate - # the route, but a direct backend caller (bench, plugin, test) would otherwise start a - # worker and do checkpoint/download work before an invalid value fails deep in the load. + # Reject malformed cache-quality / cfg-parallel here too: a direct backend caller + # (bench, plugin, test) would otherwise do checkpoint/download work before failing deep. normalize_cache_quality(transformer_cache_quality) normalize_cfg_parallel(cfg_parallel) _ensure_mp4_encoder_available() @@ -1590,8 +1588,7 @@ class VideoBackend: # THROUGH the proxy so the replica carries the same hooks and each branch's # cache state matches the single-GPU run (the bit-identity precondition). # If the primary-only cache cannot be removed, reapplying through the proxy - # would double-hook the primary and desync the branches, so fail the load - # (the _precommit_cfg_parallel rollback then tears the proxy back down). + # double-hooks the primary and desyncs the branches, so fail the load. if not _disengage_step_cache( cfg_parallel_proxy._primary, reason = "re-engaging through the cfg-parallel proxy",