Gate explicit attention kernels on NVIDIA CUDA and roll back partial FBCache hooks

studio/backend/core/inference/diffusion_attention.py

@@ -129,6 +129,11 @@ def select_attention_backend(
         backend = _ALIASES[alias]
         if backend == "native":
             return None
+        # Every explicit kernel here (cuDNN / flash* / sage) is CUDA+NVIDIA-only; on
+        # ROCm / MPS / CPU diffusers accepts the name at set time and the first
+        # generation crashes, so drop to the native default up front.
+        if not _is_cuda_nvidia(target):
+            return None
         # An arch-gated kernel (flash3/flash4) on a card that can't run it would set fine
         # then crash mid-generation, so drop it to the native default up front.
         if not _backend_arch_supported(backend):

studio/backend/core/inference/diffusion_cache.py

@@ -89,7 +89,10 @@ def apply_step_cache(
         _warn(logger, mode, RuntimeError("transformer has no cache_context (not a CacheMixin)"))
         return None
     try:
-        from diffusers import FirstBlockCacheConfig
+        try:
+            from diffusers import FirstBlockCacheConfig
+        except ImportError:  # older diffusers exports it only from diffusers.hooks
+            from diffusers.hooks import FirstBlockCacheConfig
 
         config = FirstBlockCacheConfig(threshold = thr)
         enable_cache(config)
@@ -101,6 +104,12 @@ def apply_step_cache(
             logger.info("diffusion.cache: %s engaged (threshold=%s)", mode, thr)
         return mode
     except Exception as exc:  # noqa: BLE001 — incompatible model -> run uncached
+        # enable_cache can fail after hooking some blocks; drop any partial hooks so
+        # the reported-uncached model doesn't actually run half-cached.
+        try:
+            transformer.disable_cache()
+        except Exception:  # noqa: BLE001
+            pass
         _warn(logger, mode, exc)
         return None
 

studio/backend/tests/test_diffusion_attention.py

@@ -75,7 +75,7 @@ def test_auto_stays_native_off_nvidia(monkeypatch):
 
 
 def test_explicit_backend_honored_regardless_of_speed(monkeypatch):
-    monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False)
+    monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True)
     # Pin a high capability so the arch-gated flash4 isn't dropped by the runtime check.
     monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0))
     assert select_attention_backend(_target(), "sage", speed_active = False) == "sage"
@@ -83,6 +83,15 @@ def test_explicit_backend_honored_regardless_of_speed(monkeypatch):
     assert select_attention_backend(_target(), "cudnn", speed_active = False) == "_native_cudnn"
 
 
+def test_explicit_backend_dropped_off_nvidia_cuda(monkeypatch):
+    # Explicit cuDNN/flash/sage on ROCm / MPS / CPU passes diffusers' set-time check
+    # and crashes at the first generation, so selection drops to the native default.
+    monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False)
+    monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0))
+    for alias in ("sage", "flash", "flash4", "cudnn"):
+        assert select_attention_backend(_target(device = "mps"), alias, speed_active = True) is None
+
+
 def test_explicit_native_returns_none():
     # native is the default -> nothing to set.
     assert select_attention_backend(_target(), "native", speed_active = True) is None

studio/backend/tests/test_diffusion_cache.py

@@ -136,6 +136,29 @@ def test_incompatible_model_runs_uncached(monkeypatch):
     assert apply_step_cache(_pipe(t), mode = "fbcache") is None
 
 
+def test_enable_cache_failure_rolls_back_partial_hooks(monkeypatch):
+    # enable_cache can raise after hooking some blocks; the reported-uncached model
+    # must not actually run half-cached, so the failure path calls disable_cache.
+    _stub_diffusers(monkeypatch)
+    t = _MixinTransformer(fail = True)
+    t.disabled = False
+    t.disable_cache = lambda: setattr(t, "disabled", True)
+    assert apply_step_cache(_pipe(t), mode = "fbcache") is None
+    assert t.disabled is True
+
+
+def test_config_import_falls_back_to_hooks_module(monkeypatch):
+    # Older diffusers exports FirstBlockCacheConfig only from diffusers.hooks.
+    diffusers = types.ModuleType("diffusers")  # no FirstBlockCacheConfig attribute
+    monkeypatch.setitem(sys.modules, "diffusers", diffusers)
+    hooks = types.ModuleType("diffusers.hooks")
+    hooks.FirstBlockCacheConfig = _Config
+    monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
+    t = _MixinTransformer()
+    assert apply_step_cache(_pipe(t), mode = "fbcache") == TC_FBCACHE
+    assert t.enabled_with.threshold == DEFAULT_FBCACHE_THRESHOLD
+
+
 def test_missing_transformer_is_none(monkeypatch):
     _stub_diffusers(monkeypatch)
     pipe = types.SimpleNamespace(transformer = None)