From 76eee534eabd8df6d4b05e3dacc0c9083f1d09a5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 01:48:12 +0000 Subject: [PATCH] Gate explicit attention kernels on NVIDIA CUDA and roll back partial FBCache hooks --- .../core/inference/diffusion_attention.py | 5 ++++ .../backend/core/inference/diffusion_cache.py | 11 ++++++++- .../backend/tests/test_diffusion_attention.py | 11 ++++++++- studio/backend/tests/test_diffusion_cache.py | 23 +++++++++++++++++++ 4 files changed, 48 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index e652b0068c..8b1ce29ec8 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/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): diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index b7a2b7ac45..3cb74d7f80 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/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 diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 3e43aab8a9..99c5e966f2 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/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 diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 62071d9aa6..fe2f331281 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/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)