diff --git a/scripts/fp8_layer_ablation.py b/scripts/fp8_layer_ablation.py index 9374c8fb45..71e52ebd60 100644 --- a/scripts/fp8_layer_ablation.py +++ b/scripts/fp8_layer_ablation.py @@ -452,6 +452,10 @@ 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. + pipe = pipe.to("cuda") print(f"[load] pipe built in {time.perf_counter()-t0:.1f}s", flush = True) tup = _capture_forward_tuple( diff --git a/scripts/sdpa_mask_backend_probe.py b/scripts/sdpa_mask_backend_probe.py index 4f2a030869..f9d3f71c63 100644 --- a/scripts/sdpa_mask_backend_probe.py +++ b/scripts/sdpa_mask_backend_probe.py @@ -19,18 +19,23 @@ def mk(): def timed(fn, iters = 20): - torch.cuda.synchronize() - for _ in range(3): - try: + try: + torch.cuda.synchronize() + for _ in range(3): fn() - except Exception as e: # noqa: BLE001 - return f"UNSUPPORTED ({type(e).__name__})" - torch.cuda.synchronize() - t0 = time.perf_counter() - for _ in range(iters): - fn() - torch.cuda.synchronize() - return (time.perf_counter() - t0) / iters * 1e3 + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn() + 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. + torch.cuda.empty_cache() + return "OOM" + except Exception as e: # noqa: BLE001 + return f"UNSUPPORTED ({type(e).__name__})" q, k, v = mk(), mk(), mk() diff --git a/scripts/video_speedmem_bench.py b/scripts/video_speedmem_bench.py index a9121aa180..7aa10b5006 100644 --- a/scripts/video_speedmem_bench.py +++ b/scripts/video_speedmem_bench.py @@ -623,11 +623,17 @@ 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 - _disengage_step_cache( + # 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. + if not _disengage_step_cache( transformer, reason = f"explicit magcache re-interpolating for {steps} steps", logger = logger, - ) + ): + raise RuntimeError( + "could not disable the existing MagCache before resizing it for " + f"{steps} steps; restart the benchmark with a fresh pipeline" + ) return apply_step_cache( view, mode = "magcache", diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 110e9c34f7..feb9a26060 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -64,6 +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_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 ebb165fc13..7375624274 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -675,10 +675,20 @@ def apply_step_cache( # reported-uncached model isn't half-cached. Restore armed compiled inners FIRST # (remove_hook splices original_forward back into module.forward). _restore_hooked_block_inners(transformer) - try: - transformer.disable_cache() - except Exception: # noqa: BLE001 - pass + disable_cache = getattr(transformer, "disable_cache", None) + if callable(disable_cache): + try: + 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. + raise RuntimeError( + "step-cache enable failed and rollback also failed; the transformer may be " + "partially cached and must be reloaded " + f"(enable error: {exc}; rollback error: {rollback_exc})" + ) from rollback_exc _warn(logger, mode, exc) return None @@ -770,10 +780,17 @@ def maybe_toggle_step_cache( and mode == TC_MAGCACHE # endswith, not substring: "#s5" would match inside "#s50". and not str(engaged).endswith(f"#s{int(steps)}") - and _disengage_step_cache( - transformer, reason = f"magcache re-interpolating for {steps} steps", logger = logger - ) ): + # 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. + if not _disengage_step_cache( + transformer, reason = f"magcache re-interpolating for {steps} steps", logger = logger + ): + raise RuntimeError( + "could not disable the existing MagCache before resizing it for " + f"{steps} steps; reload the video model before generating" + ) engaged = None if want and not engaged: return apply_step_cache( @@ -788,13 +805,18 @@ def maybe_toggle_step_cache( logger = logger, ) if not want and engaged: - if _disengage_step_cache( + # 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. + if not _disengage_step_cache( transformer, reason = f"auto: {steps} steps < {FBCACHE_MIN_STEPS}", logger = logger, ): - return None - return mode + raise RuntimeError( + "could not disable the existing step cache for a short generation; " + "reload the video model before generating" + ) + return None return mode if engaged else None diff --git a/studio/backend/core/inference/diffusion_cfg_parallel.py b/studio/backend/core/inference/diffusion_cfg_parallel.py index 8c073851c8..58485409e7 100644 --- a/studio/backend/core/inference/diffusion_cfg_parallel.py +++ b/studio/backend/core/inference/diffusion_cfg_parallel.py @@ -230,13 +230,24 @@ class CFGParallelProxy: raise def disable_cache(self) -> None: - self._primary.disable_cache() - try: - self._replica.disable_cache() - self._broken = False - except Exception as exc: # noqa: BLE001 -- replica out of sync: stop routing - self._broken = True - _warn(self._logger, "cfg-parallel replica disable_cache", exc) + # 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. + failures: list[tuple[str, Exception]] = [] + for name, module in (("primary", self._primary), ("replica", self._replica)): + try: + module.disable_cache() + except Exception as exc: # noqa: BLE001 -- collect, don't skip the other branch + failures.append((name, exc)) + _warn(self._logger, f"cfg-parallel {name} disable_cache", exc) + self._broken = bool(failures) + if failures: + details = "; ".join(f"{name}: {exc}" for name, exc in failures) + raise RuntimeError( + f"CFG-parallel cache removal failed ({details}); parallel routing is " + "disabled until the model is reloaded" + ) def _reset_stateful_cache(self, *args: Any, **kwargs: Any) -> None: for module in (self._primary, self._replica): @@ -401,6 +412,14 @@ 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. + 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") + ) q, k, v = (x.permute(0, 2, 1, 3).contiguous() for x in (query, key, value)) out = torch.ops.aten._scaled_dot_product_cudnn_attention( q, k, v, attn_mask, False, 0.0, is_causal, False, scale = scale diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index d0880d9430..99350a052d 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -451,9 +451,32 @@ def _step_cache_all_or_none( expert(s) and report the cache off. Otherwise the whole MoE reports cached while half the schedule runs uncached. Returns (mode-or-None, failure-reason-or-None); a single-DiT family can never see a mixed outcome.""" + pairs = list(zip(_views_for(pipe, fam), _transformer_names(pipe, fam))) results: list[tuple[Any, str, Optional[str]]] = [] - for view, expert_name in zip(_views_for(pipe, fam), _transformer_names(pipe, fam)): - results.append((view, expert_name, engage_fn(view, expert_name))) + try: + 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). + rollback_failed: list[str] = [] + for view, name in pairs: + transformer = getattr(view, "transformer", None) + if getattr(transformer, "_unsloth_step_cache", None) is None: + continue + if not _disengage_step_cache( + transformer, reason = "all-or-none transaction aborted", logger = logger + ): + rollback_failed.append(name) + if rollback_failed: + raise RuntimeError( + "step-cache transaction raised and rollback failed for " + + ", ".join(rollback_failed) + + "; reload the video model before generating" + ) from exc + raise engaged = [(view, name, mode) for view, name, mode in results if mode is not None] if engaged and len(engaged) < len(results): missing = ", ".join(name for _, name, mode in results if mode is None) @@ -518,6 +541,8 @@ class VideoBackend: transformer_quant: Optional[str] = None, text_encoder_quant: Optional[str] = None, vae_quant: Optional[str] = None, + transformer_cache_quality: Optional[str] = None, + cfg_parallel: Optional[str] = None, ) -> VideoFamily: """Cheap, network-free validation shared by the route and the load path.""" kind = resolve_video_model_kind(gguf_filename, model_kind) @@ -620,6 +645,11 @@ 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. + normalize_cache_quality(transformer_cache_quality) + normalize_cfg_parallel(cfg_parallel) _ensure_mp4_encoder_available() return fam @@ -656,6 +686,8 @@ class VideoBackend: transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant, vae_quant = vae_quant, + transformer_cache_quality = transformer_cache_quality, + cfg_parallel = cfg_parallel, ) with self._lock: if self._loading is not None and self._loading.error is None: @@ -1039,6 +1071,8 @@ class VideoBackend: transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant, vae_quant = vae_quant, + transformer_cache_quality = transformer_cache_quality, + cfg_parallel = cfg_parallel, ) kind = resolve_video_model_kind(gguf_filename, model_kind) # An explicit Speed="off" (bit-exact) load pins the companions dense too: promoting @@ -1555,11 +1589,18 @@ class VideoBackend: # The cache engaged on the primary BEFORE the proxy existed; re-engage # 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). - _disengage_step_cache( + # 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). + if not _disengage_step_cache( cfg_parallel_proxy._primary, reason = "re-engaging through the cfg-parallel proxy", logger = logger, - ) + ): + raise RuntimeError( + "could not disable the primary-only step cache before installing " + "CFG-parallel cache hooks; reload the video model before generating" + ) cache_engaged = apply_step_cache( pipe, mode = cache_request, diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 6f84cfd66a..26dfdd8f38 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -107,6 +107,8 @@ async def load_video_model( transformer_quant = request.transformer_quant, text_encoder_quant = request.text_encoder_quant, vae_quant = request.vae_quant, + transformer_cache_quality = request.transformer_cache_quality, + cfg_parallel = request.cfg_parallel, ) # Refuse while training is running (VRAM competition). Mirrors the image-load guard. _guard_video_load_against_training() diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index a396b2628c..5bd357d135 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -143,6 +143,16 @@ def test_flash3_dropped_on_blackwell(monkeypatch): assert select_attention_backend(_target(), "flash3", speed_active = False) == "_flash_3_hub" +def test_flash2_dropped_below_ampere(monkeypatch): + # Dao-AILab FlashAttention 2 needs Ampere (SM80)+; on Turing (SM75) diffusers accepts the + # backend then crashes at generation, so the arch gate must drop it like cuDNN/FA3/FA4. + monkeypatch.setattr(att, "_cuda_capability", lambda: (7, 5)) + assert select_attention_backend(_target(), "flash", speed_active = False) is None + # Ampere+ still honors it. + monkeypatch.setattr(att, "_cuda_capability", lambda: (8, 0)) + assert select_attention_backend(_target(), "flash", speed_active = False) == "flash" + + def test_explicit_cudnn_dropped_below_sm80(monkeypatch): # An explicit cuDNN request on pre-Ampere (T4 SM75 / V100 SM70) must drop to native, # not set fine and crash at first generation -- the same gate the auto path applies. @@ -551,6 +561,13 @@ def test_backend_supported_on_device_cudnn_needs_ampere(monkeypatch): assert att.attention_backend_supported_on_device("_native_cudnn", 1) is False +def test_backend_supported_on_device_flash2_needs_ampere(monkeypatch): + # FlashAttention 2 needs Ampere+ (SM80): rejected on a pre-Ampere (T4/SM75) replica. + _stub_cuda_capability(monkeypatch, {0: (8, 0), 1: (7, 0)}) + assert att.attention_backend_supported_on_device("flash", 0) is True + assert att.attention_backend_supported_on_device("flash", 1) is False + + def test_backend_supported_on_device_unqueryable_is_permissive(monkeypatch): # An unqueryable device must not block on a guess (best-effort, like _backend_arch_supported). torch = types.ModuleType("torch") diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 38c18d3c5a..dafd396b3b 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -372,6 +372,49 @@ def test_toggle_reengages_after_a_disable(monkeypatch): assert mode == TC_FBCACHE and t.enables == 2 +def test_toggle_magcache_step_change_hard_errors_when_disable_fails(monkeypatch): + # A failed removal during a magcache step-count change used to fall through and report + # "magcache" while the OLD #sN curve stayed armed (wrong ratio schedule). Fail closed. + import core.inference.diffusion_cache as dc_mod + + _stub_diffusers(monkeypatch) + t = _ToggleTransformer() + t._unsloth_step_cache = "magcache@0.06#s50" + monkeypatch.setattr(dc_mod, "_disengage_step_cache", lambda *a, **k: False) + with pytest.raises(RuntimeError, match = "reload the video model"): + maybe_toggle_step_cache( + _pipe(t), steps = 30, mode = dc_mod.TC_MAGCACHE, family = "hunyuanvideo-1.5" + ) + + +def test_toggle_below_bar_hard_errors_when_disable_fails(monkeypatch): + # Below the cache threshold we want uncached; a failed disable used to report the stale + # mode instead of surfacing that the (wrong-step) cache is still armed. + import core.inference.diffusion_cache as dc_mod + + _stub_diffusers(monkeypatch) + t = _ToggleTransformer() + t._unsloth_step_cache = "magcache@0.06#s50" + monkeypatch.setattr(dc_mod, "_disengage_step_cache", lambda *a, **k: False) + with pytest.raises(RuntimeError, match = "reload the video model"): + maybe_toggle_step_cache( + _pipe(t), + steps = FBCACHE_MIN_STEPS - 1, + mode = dc_mod.TC_MAGCACHE, + family = "hunyuanvideo-1.5", + ) + + +def test_apply_step_cache_enable_and_cleanup_failure_requires_reload(monkeypatch): + # enable_cache fails AFTER partially hooking and the cleanup disable_cache ALSO fails: + # the transformer may keep partial hooks, so surface it instead of a clean uncached None. + _stub_diffusers(monkeypatch) + t = _MixinTransformer(fail = True) + t.disable_cache = lambda: (_ for _ in ()).throw(RuntimeError("cleanup failed")) + with pytest.raises(RuntimeError, match = "partially cached and must be reloaded"): + apply_step_cache(_pipe(t), mode = "fbcache") + + def test_toggle_noop_without_cache_support(monkeypatch): _stub_diffusers(monkeypatch) t = _NonCacheMixinTransformer() diff --git a/studio/backend/tests/test_diffusion_cfg_parallel.py b/studio/backend/tests/test_diffusion_cfg_parallel.py index 231f859060..da1b8235b5 100644 --- a/studio/backend/tests/test_diffusion_cfg_parallel.py +++ b/studio/backend/tests/test_diffusion_cfg_parallel.py @@ -411,6 +411,26 @@ def test_replica_enable_failure_reraises_and_breaks(monkeypatch): proxy.shutdown() +def test_disable_cache_primary_failure_still_cleans_replica_and_breaks(monkeypatch): + # The primary's disable_cache used to run OUTSIDE the guard: if it raised, the replica was + # never cleaned and _broken stayed False, so a half-removed pair kept routing. Removal must + # be transactional -- clean both branches, mark broken, and surface a reload-required error. + proxy, primary, replica, _ = _make_proxy(monkeypatch) + + def _boom(): + raise RuntimeError("primary hook removal failed") + + primary.disable_cache = _boom + with pytest.raises(RuntimeError, match = "CFG-parallel cache removal failed"): + proxy.disable_cache() + assert proxy._broken is True + assert replica.disables == 1 + # A broken proxy pins the sequential passthrough (no parallel routing). + plan = proxy.plan_generation(cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33) + assert plan["enabled"] is False + proxy.shutdown() + + def test_reset_stateful_cache_fans_out(monkeypatch): proxy, primary, replica, _ = _make_proxy(monkeypatch) proxy._reset_stateful_cache() diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index d996922e9f..454fd6b1ec 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -2159,6 +2159,30 @@ def test_step_cache_all_or_none_rolls_back_first_expert_failure(monkeypatch): assert disengaged == [t2] +def test_step_cache_all_or_none_rolls_back_when_later_expert_raises(monkeypatch): + # A later expert RAISING mid-loop (not returning None) must not leave earlier experts + # engaged: the helper disengages any expert that got a cache marker, then re-raises. + import core.inference.video as video + + pipe, fam, t1, _t2 = _moe_pipe_and_fam() + disengaged: list = [] + monkeypatch.setattr( + video, + "_disengage_step_cache", + lambda transformer, *, reason, logger = None: disengaged.append(transformer) or True, + ) + + def engage(view, expert_name): + if expert_name == "transformer": + view.transformer._unsloth_step_cache = "magcache@0.1#s30" + return "magcache" + raise RuntimeError("expert 2 boom") + + with pytest.raises(RuntimeError, match = "expert 2 boom"): + video._step_cache_all_or_none(pipe, fam, engage, logger = None) + assert disengaged == [t1] + + def test_step_cache_all_or_none_uniform_outcomes(monkeypatch): # Both experts engaged -> the mode is reported with no rollback; neither engaged # -> plain uncached with no failure reason (the pre-existing best-effort path). diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 63e92beaf4..e4602a658d 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -104,6 +104,8 @@ class _FakeBackend(video_module.VideoBackend): transformer_quant = None, text_encoder_quant = None, vae_quant = None, + transformer_cache_quality = None, + cfg_parallel = None, ): # Mirror the real backend's cheap validation so the route's # validate-before-evict ordering is exercised.