Studio: harden video/diffusion cache, attention, and CFG-parallel fault paths

- diffusion_attention: arch-gate FlashAttention 2 to Ampere (SM80)+ in both the
  primary selector and the heterogeneous-replica guard (it crashed on pre-Ampere).
- diffusion_cfg_parallel: convert boolean attn masks to additive bias before the direct
  cuDNN op so partial masks match F.scaled_dot_product_attention; make proxy disable_cache
  transactional (clean both branches, mark broken, surface a reload-required error).
- diffusion_cache: fail closed when a magcache step-count resize or below-threshold
  disable cannot remove the old cache; surface a failed enable+cleanup instead of a false
  uncached None.
- video: roll back earlier experts when a later expert raises in the all-or-none step-cache
  loop; fail the load when the primary-only cache cannot be re-engaged through the
  CFG-parallel proxy; validate transformer_cache_quality and cfg_parallel before the worker.
- scripts: place the fp8 ablation pipeline on CUDA; fail closed on a failed magcache resize
  in the speedmem bench; label OOM distinctly in the SDPA mask probe.
- tests: regressions for the FA2 arch gate, transactional proxy disable, all-or-none
  exception rollback, magcache fail-closed transitions, and enable+cleanup failure.

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(

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()

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",

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)+
 }

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
 
 

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

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,

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()

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")

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()

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()

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).

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.