From 7bd15727c28033123411f5aefc52268539b6c5f4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 29 Jun 2026 10:43:04 +0000 Subject: [PATCH] Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. --- .../core/inference/diffusion_attention.py | 64 ++++++++++++++----- .../backend/tests/test_diffusion_attention.py | 49 ++++++++++++++ 2 files changed, 98 insertions(+), 15 deletions(-) diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index b60e737ffc..04f812a13e 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -68,11 +68,13 @@ def normalize_attention_backend(value: Optional[str]) -> Optional[str]: # Backends diffusers validates only by *package* at set time (``_check_attention_backend_ # requirements`` checks the ``kernels`` install, not the GPU), but whose kernels need a # specific CUDA arch at run time -- so an explicit request on the wrong card loads/sets fine -# and then crashes mid-generation. Gate them up front: minimum (major, minor) compute -# capability per dispatcher backend name. -_MIN_CUDA_CAPABILITY: dict[str, tuple[int, int]] = { - "_flash_3_hub": (9, 0), # FlashAttention 3 -> Hopper (SM90) - "flash_4_hub": (10, 0), # FlashAttention 4 -> Blackwell (SM100) +# and then crashes mid-generation. Gate them up front by a (min, max-exclusive) compute +# capability range. FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel, so it +# needs an upper bound: an explicit flash3 on a B200 (SM100) must drop to native instead of +# setting fine then crashing at generation. FlashAttention 4 is Blackwell+ (no upper bound). +_ARCH_CAPABILITY: dict[str, tuple[tuple[int, int], Optional[tuple[int, int]]]] = { + "_flash_3_hub": ((9, 0), (10, 0)), # FlashAttention 3 -> Hopper (SM90) only + "flash_4_hub": ((10, 0), None), # FlashAttention 4 -> Blackwell (SM100)+ } @@ -88,18 +90,19 @@ def _cuda_capability() -> Optional[tuple[int, int]]: def _backend_arch_supported(backend: str) -> bool: - """False only when ``backend`` needs a known-higher CUDA arch than this device has. + """False only when ``backend`` needs a CUDA arch outside this device's supported range. Unknown capability (no CUDA / detection failure) returns True so we never block on a guess -- diffusers' own set-time check still guards the package, and a genuine run-time failure falls back to native.""" - required = _MIN_CUDA_CAPABILITY.get(backend) - if required is None: + bounds = _ARCH_CAPABILITY.get(backend) + if bounds is None: return True have = _cuda_capability() if have is None: return True - return have >= required + low, high = bounds + return have >= low and (high is None or have < high) def _is_cuda_nvidia(target: Any) -> bool: @@ -130,6 +133,11 @@ def select_attention_backend( # then crash mid-generation, so drop it to the native default up front. if not _backend_arch_supported(backend): return None + # cuDNN fused SDPA needs Ampere+ (SM80); diffusers accepts it on pre-SM80 cards + # (T4/V100) then fails at the first generation, so apply the same gate to an + # explicit cuDNN request as the auto path already does. + if backend == "_native_cudnn" and not _cudnn_attention_supported(): + return None return backend # auto if speed_active and _is_cuda_nvidia(target) and _cudnn_attention_supported(): @@ -170,6 +178,12 @@ def apply_attention_backend( if backend is not None: try: fn(backend) + # set_attention_backend also pins the backend in diffusers' process-wide + # registry. This transformer's own processors keep it locally (their + # _attention_backend is now explicit), so reset the global default back to + # native -- otherwise a later component whose processors are unconfigured + # (backend None) silently inherits this kernel. + _reset_global_backend_to_native(logger) if logger is not None: logger.info("diffusion.attention: backend=%s", backend) return backend @@ -186,17 +200,37 @@ def _active_attention_backend() -> Optional[str]: try: from diffusers.models.attention_dispatch import _AttentionBackendRegistry - # get_active_backend() returns an AttentionBackend enum (or None), NOT a tuple: - # unpacking it as `name, _` raises ValueError (swallowed below), so this always - # returned None and the native-restore short-circuit never fired. - backend = _AttentionBackendRegistry.get_active_backend() - if backend is None: + # get_active_backend() returns a (AttentionBackendName, fn) tuple (or None), so + # take element 0 and read its .value (e.g. "native"); reading .value off the + # tuple itself would yield a junk string that never compares equal to a name. + active = _AttentionBackendRegistry.get_active_backend() + if active is None: return None - return getattr(backend, "value", str(backend)) + name = active[0] if isinstance(active, tuple) else active + return getattr(name, "value", str(name)) except Exception: # noqa: BLE001 return None +def _reset_global_backend_to_native(logger: Any) -> None: + """Reset diffusers' process-wide active attention backend to native after a + successful per-transformer set, so a later component whose processors are + unconfigured (backend None) does not inherit this transformer's kernel. The + transformer's own processors keep the backend just set. Best-effort and silent: + if the diffusers internals move, the prior (leaking) behavior is unchanged.""" + if _active_attention_backend() == ATTN_NATIVE: + return + try: + from diffusers.models.attention_dispatch import ( + AttentionBackendName, + _AttentionBackendRegistry, + ) + + _AttentionBackendRegistry.set_active_backend(AttentionBackendName.NATIVE) + except Exception: # noqa: BLE001 — best-effort; leave the global as-is on any change + pass + + def _restore_native_backend(set_backend_fn: Any, logger: Any) -> None: """Force the native default when the global active backend isn't already native.""" if _active_attention_backend() == ATTN_NATIVE: diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 2c9977b624..3e43aab8a9 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -107,6 +107,28 @@ def test_arch_gate_does_not_block_when_capability_unknown(monkeypatch): assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub" +def test_flash3_dropped_on_blackwell(monkeypatch): + # FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel: an explicit + # flash3 on a B200 (SM100) must drop to native rather than set fine then crash. + monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0)) + assert select_attention_backend(_target(), "flash3", speed_active = False) is None + # FA4 is still honored on Blackwell. + assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub" + # flash3 is allowed exactly on Hopper SM90. + monkeypatch.setattr(att, "_cuda_capability", lambda: (9, 0)) + assert select_attention_backend(_target(), "flash3", speed_active = False) == "_flash_3_hub" + + +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. + monkeypatch.setattr(att, "_cuda_capability", lambda: (7, 5)) + assert select_attention_backend(_target(), "cudnn", speed_active = False) is None + # Ampere+ still honors it. + monkeypatch.setattr(att, "_cuda_capability", lambda: (8, 0)) + assert select_attention_backend(_target(), "cudnn", speed_active = False) == "_native_cudnn" + + # ── apply ───────────────────────────────────────────────────────────────────────── class _FakeTransformer: def __init__(self, *, fail = False): @@ -174,3 +196,30 @@ def test_apply_failed_kernel_restores_native_when_polluted(monkeypatch): def test_apply_handles_missing_method(): pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) assert apply_attention_backend(pipe, "_native_cudnn") is None + + +def test_apply_resets_global_registry_after_success(monkeypatch): + # After a successful per-transformer set, the process-wide registry must be reset to + # native so a later component (unconfigured processors) can't inherit this kernel -- + # while the transformer's own backend stays the engaged one. + called = {"reset": False} + monkeypatch.setattr( + att, "_reset_global_backend_to_native", lambda logger: called.__setitem__("reset", True) + ) + t = _FakeTransformer() + engaged = apply_attention_backend(_pipe(t), "_native_cudnn") + assert engaged == "_native_cudnn" and t.set_to == "_native_cudnn" + assert called["reset"] is True + + +def test_active_attention_backend_reads_tuple_return(): + # get_active_backend() returns a (AttentionBackendName, fn) tuple; the helper must read + # the name's .value, not stringify the tuple (which never compares equal to a name). + pytest.importorskip("diffusers") + from diffusers.models.attention_dispatch import ( + AttentionBackendName, + _AttentionBackendRegistry, + ) + + _AttentionBackendRegistry.set_active_backend(AttentionBackendName.NATIVE) + assert att._active_attention_backend() == "native"