Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa

- apply_attention_backend now restores the native default when no backend is requested or a
  kernel fails. diffusers keeps a process-wide active attention backend that
  set_attention_backend updates, and a fresh transformer's processors follow it, so a load
  that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile
  load pinned, breaking the bit-identical/off guarantee.
- select_attention_backend drops flash3/flash4 up front when the CUDA capability is below
  Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit
  request on the wrong card set fine then crashed mid-generation; it now falls back to native.
- Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a
  valid alias of native) is accepted instead of 422-rejected by Pydantic.
- Drop the dead replace('-','_') normalization (no alias uses dashes/underscores).
- perf_levers_probe.py output dir is now relative to the script, not a hardcoded path.

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}

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)

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

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