diff --git a/scripts/perf_levers_probe.py b/scripts/perf_levers_probe.py index 92f2b89fdc..e0ea8ba312 100644 --- a/scripts/perf_levers_probe.py +++ b/scripts/perf_levers_probe.py @@ -25,7 +25,7 @@ import numpy as np BASE = "Tongyi-MAI/Z-Image-Turbo" PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" -OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/perf_levers_images") +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "perf_levers_images" _LP = {"fn": None} diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index cf74b0fb0b..eea5e91623 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -55,7 +55,7 @@ def normalize_attention_backend(value: Optional[str]) -> Optional[str]: Raises ValueError for an unsupported alias so a bad request is rejected cheaply.""" if value is None: return ATTN_AUTO - normalized = str(value).strip().lower().replace("-", "_") + normalized = str(value).strip().lower() if not normalized: return ATTN_AUTO if normalized not in ATTN_ALIASES: @@ -65,6 +65,43 @@ def normalize_attention_backend(value: Optional[str]) -> Optional[str]: return normalized +# 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) +} + + +def _cuda_capability() -> Optional[tuple[int, int]]: + """(major, minor) compute capability of the active CUDA device, or None if unknown.""" + try: + import torch + if not torch.cuda.is_available(): + return None + return tuple(torch.cuda.get_device_capability()) # type: ignore[return-value] + except Exception: # noqa: BLE001 + return None + + +def _backend_arch_supported(backend: str) -> bool: + """False only when ``backend`` needs a known-higher CUDA arch than this device has. + + 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: + return True + have = _cuda_capability() + if have is None: + return True + return have >= required + + def _is_cuda_nvidia(target: Any) -> bool: """CUDA device on an NVIDIA (non-ROCm) build -- where cuDNN attention applies.""" if getattr(target, "device", None) != "cuda": @@ -87,7 +124,13 @@ def select_attention_backend( alias = normalize_attention_backend(requested) if alias != ATTN_AUTO: backend = _ALIASES[alias] - return None if backend == "native" else backend + if backend == "native": + 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): + return None + return backend # auto if speed_active and _is_cuda_nvidia(target): return "_native_cudnn" @@ -102,25 +145,53 @@ def apply_attention_backend( ) -> Optional[str]: """Set ``backend`` on ``pipe.transformer`` via the diffusers dispatcher. - Returns the backend actually engaged, or None when left at the default (either because - ``backend`` was None or because the requested kernel was unavailable -> graceful - fallback to the diffusers default, never a load failure). Best-effort.""" - if backend is None: - return None + Returns the backend actually engaged, or None when left at the native default (either + because ``backend`` was None or because the requested kernel was unavailable -> graceful + fallback, never a load failure). + + diffusers keeps a *process-wide* active attention backend that ``set_attention_backend`` + also updates, and a fresh transformer's processors follow it (their ``_attention_backend`` + defaults to None). So a load that wants native must restore it explicitly: otherwise it + silently inherits a backend an earlier load pinned (e.g. cuDNN under a speed profile), + breaking the bit-identical/``off`` guarantee. Best-effort throughout.""" transformer = getattr(pipe, "transformer", None) fn = getattr(transformer, "set_attention_backend", None) if not callable(fn): return None + if backend is not None: + try: + fn(backend) + if logger is not None: + logger.info("diffusion.attention: backend=%s", backend) + return backend + except Exception as exc: # noqa: BLE001 — unavailable kernel -> restore native below + _warn(logger, backend, exc) + # No backend requested, or the requested one failed: pin the native default so a stale + # process-wide backend from a previous load can't leak into this one. + _restore_native_backend(fn, logger) + return None + + +def _active_attention_backend() -> Optional[str]: + """The diffusers process-wide active attention backend name, or None if undeterminable.""" try: - fn(backend) - if logger is not None: - logger.info("diffusion.attention: backend=%s", backend) - return backend - except Exception as exc: # noqa: BLE001 — unavailable kernel -> diffusers default - _warn(logger, backend, exc) + from diffusers.models.attention_dispatch import _AttentionBackendRegistry + name, _ = _AttentionBackendRegistry.get_active_backend() + return getattr(name, "value", str(name)) + except Exception: # noqa: BLE001 return None +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: + return # already native -> avoid redundant work and an extra dispatcher warning + try: + set_backend_fn(ATTN_NATIVE) + except Exception as exc: # noqa: BLE001 — best-effort restore + _warn(logger, ATTN_NATIVE, exc) + + def _warn(logger: Any, what: str, exc: Exception) -> None: if logger is not None: logger.warning("diffusion.attention: %s unavailable (%s); using default", what, exc) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index d47ad240dd..ddfdd2ac7c 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1750,6 +1750,7 @@ class DiffusionLoadRequest(BaseModel): Literal[ "auto", "native", + "sdpa", "cudnn", "flash", "flash2", @@ -1764,7 +1765,7 @@ class DiffusionLoadRequest(BaseModel): description = "Attention kernel via the diffusers dispatcher. auto picks the best " "exact backend for the device (cuDNN fused attention on NVIDIA, ~1.18x and " "near-lossless, when a speed profile is active; native SDPA elsewhere and when " - "speed=off). native forces default SDPA; cudnn/flash/flash3/flash4 are exact " + "speed=off). native (alias sdpa) forces default SDPA; cudnn/flash/flash3/flash4 are exact " "(kernel/arch-gated); sage is INT8 attention (a small quality cost, consumer " "friendly); xformers/aiter are memory-efficient (NVIDIA) / AMD ROCm. An " "unavailable kernel falls back to the default.", diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index bb006088c2..2c5c94b06f 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -32,11 +32,20 @@ def test_normalize_defaults_and_aliases(): assert normalize_attention_backend("auto") == ATTN_AUTO assert normalize_attention_backend("CuDNN") == "cudnn" assert normalize_attention_backend("FLASH3") == "flash3" + assert normalize_attention_backend("sdpa") == "sdpa" def test_normalize_rejects_unknown(): with pytest.raises(ValueError): normalize_attention_backend("bogus") + # dashes are no longer silently rewritten to underscores -> a dashed alias is rejected. + with pytest.raises(ValueError): + normalize_attention_backend("flash-3") + + +def test_sdpa_alias_maps_to_native(): + # sdpa is an alias for native -> nothing to set on the dispatcher. + assert select_attention_backend(_target(), "sdpa", speed_active = True) is None # ── select policy ───────────────────────────────────────────────────────────────── @@ -58,6 +67,8 @@ 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) + # 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" assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub" assert select_attention_backend(_target(), "cudnn", speed_active = False) == "_native_cudnn" @@ -68,6 +79,25 @@ def test_explicit_native_returns_none(): assert select_attention_backend(_target(), "native", speed_active = True) is None +# ── arch gating (flash3/flash4 need a specific CUDA capability) ───────────────────── +def test_flash3_dropped_below_hopper(monkeypatch): + monkeypatch.setattr(att, "_cuda_capability", lambda: (8, 9)) # Ada / consumer + assert select_attention_backend(_target(), "flash3", speed_active = False) is None + + +def test_flash4_dropped_below_blackwell(monkeypatch): + monkeypatch.setattr(att, "_cuda_capability", lambda: (9, 0)) # Hopper, but FA4 needs SM100 + assert select_attention_backend(_target(), "flash4", speed_active = False) is None + # flash3 still allowed on Hopper. + assert select_attention_backend(_target(), "flash3", speed_active = False) == "_flash_3_hub" + + +def test_arch_gate_does_not_block_when_capability_unknown(monkeypatch): + # Unknown capability (e.g. no CUDA) must not block -> diffusers' set-time check still guards. + monkeypatch.setattr(att, "_cuda_capability", lambda: None) + assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub" + + # ── apply ───────────────────────────────────────────────────────────────────────── class _FakeTransformer: def __init__(self, *, fail = False): @@ -84,8 +114,21 @@ def _pipe(transformer): return types.SimpleNamespace(transformer = transformer) -def test_apply_none_is_noop(): - assert apply_attention_backend(_pipe(_FakeTransformer()), None) is None +def test_apply_none_leaves_native_when_global_already_native(monkeypatch): + # Global already native -> no redundant set call, returns None. + monkeypatch.setattr(att, "_active_attention_backend", lambda: "native") + t = _FakeTransformer() + assert apply_attention_backend(_pipe(t), None) is None + assert t.set_to is None + + +def test_apply_none_restores_native_when_global_polluted(monkeypatch): + # A previous load pinned cuDNN process-wide; a native load must reset it so it can't + # silently inherit cuDNN (the bit-identical/off guarantee). + monkeypatch.setattr(att, "_active_attention_backend", lambda: "_native_cudnn") + t = _FakeTransformer() + assert apply_attention_backend(_pipe(t), None) is None + assert t.set_to == "native" def test_apply_sets_backend(): @@ -94,12 +137,31 @@ def test_apply_sets_backend(): assert engaged == "_native_cudnn" and t.set_to == "_native_cudnn" -def test_apply_falls_back_on_unavailable_kernel(): +def test_apply_falls_back_on_unavailable_kernel(monkeypatch): # an unavailable kernel must not fail the load -> returns None (diffusers default). + monkeypatch.setattr(att, "_active_attention_backend", lambda: "native") t = _FakeTransformer(fail = True) assert apply_attention_backend(_pipe(t), "sage") is None +def test_apply_failed_kernel_restores_native_when_polluted(monkeypatch): + # Requested kernel fails AND the global is polluted: restore native before returning. + monkeypatch.setattr(att, "_active_attention_backend", lambda: "_native_cudnn") + + class _FailOnceTransformer: + def __init__(self): + self.calls = [] + + def set_attention_backend(self, name): + self.calls.append(name) + if name != "native": + raise RuntimeError(f"{name} kernel unavailable") + + t = _FailOnceTransformer() + assert apply_attention_backend(_pipe(t), "sage") is None + assert t.calls == ["sage", "native"] + + def test_apply_handles_missing_method(): pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) assert apply_attention_backend(pipe, "_native_cudnn") is None