fix: guide scenes crash on Blackwell — route masked attention to SDPA when xformers can't

On compute capability > (9, 0) GPUs (Blackwell sm_120, e.g. RTX PRO 6000),
xformers' memory_efficient_attention has no kernel that accepts a tensor
attn_bias — fa3F/cutlassF require cap <= (9,0), fa2F/fa3F reject a tensor
bias. So unmasked attention works (i2v ANCHOR scenes generate) but MASKED
attention crashes: scenes generated from GUIDES hit the LTX guide path
(comfy/ldm/lightricks/model.py _attention_with_guide_mask) which passes a
tensor mask → "No operator found for memory_efficient_attention_forward".

This is the ComfyUI-wide issue normally worked around with
--disable-xformers --use-pytorch-cross-attention; when that flag isn't in
effect on an instance, the guide path still reaches xformers.

Fix: _funpack_install_mask_safe_attention(model), called at the top of the
chain sampler's sample(). When the active backend is xformers AND the device
capability is > (9,0), install ComfyUI's per-call optimized_attention_override
(the wrap_attn hook) so MASKED attention falls back to attention_pytorch
(SDPA handles masks on every GPU) while UNMASKED attention stays on the fast
xformers path. Threaded via transformer_options so it reaches every scene's
forward. No global monkeypatch; no recursion (wrap_attn stamps
_inside_attn_wrapper before calling the override). No-op on supported GPUs,
non-xformers backends, or when another override is already installed;
idempotent.

Tests: tests/test_mask_safe_attention.py — 9 cases covering override routing
(masked→pytorch via kwarg or positional, unmasked→original) and install
gating (xformers+Blackwell installs; supported cap / non-xformers / no-CUDA
skip; existing override respected; idempotent). Run with the ComfyUI venv.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

tests/test_mask_safe_attention.py

@@ -0,0 +1,138 @@
+"""Unit tests for the Blackwell (sm_120) xformers masked-attention fallback in samplers.py.
+
+On GPUs newer than capability (9, 0) xformers has no kernel that accepts a tensor attn_bias,
+so i2v ANCHOR scenes (unmasked attention) generate fine but GUIDE scenes (LTX guide path
+passes a tensor mask) crash. _funpack_install_mask_safe_attention installs a per-call override
+(ComfyUI wrap_attn) that routes only MASKED attention to SDPA when that exact combo is present.
+
+Covers: override routing (masked→pytorch, unmasked→original, mask as kwarg or positional),
+and the install gating (only xformers + too-new GPU; respects/ignores other backends, existing
+overrides, missing CUDA; idempotent).
+"""
+import sys
+import types
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+# Minimal comfy stubs so `import samplers` works without a full ComfyUI env.
+for _name in (
+    "comfy", "comfy.k_diffusion", "comfy.k_diffusion.sampling",
+    "comfy.model_sampling", "comfy.nested_tensor", "comfy.sample",
+    "comfy.samplers", "comfy.utils",
+):
+    sys.modules.setdefault(_name, types.ModuleType(_name))
+sys.modules["comfy.nested_tensor"].NestedTensor = object
+
+import samplers  # noqa: E402
+
+
+def _install_attn_module(*, backend="xformers", with_pytorch_sentinel=True):
+    """Register a fake comfy.ldm.modules.attention with selectable active backend."""
+    am = types.ModuleType("comfy.ldm.modules.attention")
+
+    def attention_xformers(*a, **kw):
+        return "xformers"
+
+    def attention_pytorch(*a, **kw):
+        return "pytorch"
+
+    am.attention_xformers = attention_xformers
+    am.attention_pytorch = attention_pytorch
+    am.optimized_attention = attention_xformers if backend == "xformers" else attention_pytorch
+    for parent in ("comfy", "comfy.ldm", "comfy.ldm.modules", "comfy.ldm.modules.attention"):
+        sys.modules.setdefault(parent, types.ModuleType(parent))
+    sys.modules["comfy.ldm.modules.attention"] = am
+    return am
+
+
+def _fake_model():
+    return types.SimpleNamespace(model_options={})
+
+
+def _force_capability(monkeypatch, cap, available=True):
+    cuda = types.SimpleNamespace(
+        is_available=lambda: available,
+        get_device_capability=lambda *a, **k: cap,
+    )
+    monkeypatch.setattr(samplers.torch, "cuda", cuda, raising=False)
+
+
+# ── override routing ──────────────────────────────────────────────────────────
+def test_override_routes_masked_to_pytorch_kwarg():
+    _install_attn_module()
+    sentinel = object()
+    out = samplers._funpack_mask_safe_attention_override(
+        lambda *a, **k: "xformers", "q", "k", "v", 8, mask=sentinel,
+    )
+    assert out == "pytorch"
+
+
+def test_override_routes_masked_to_pytorch_positional():
+    _install_attn_module()
+    # (q, k, v, heads, mask) — mask as the 5th positional arg.
+    out = samplers._funpack_mask_safe_attention_override(
+        lambda *a, **k: "xformers", "q", "k", "v", 8, object(),
+    )
+    assert out == "pytorch"
+
+
+def test_override_passes_unmasked_to_original():
+    _install_attn_module()
+    out = samplers._funpack_mask_safe_attention_override(
+        lambda *a, **k: "xformers", "q", "k", "v", 8, mask=None,
+    )
+    assert out == "xformers"
+
+
+# ── install gating ────────────────────────────────────────────────────────────
+def test_installs_on_xformers_and_blackwell(monkeypatch):
+    _install_attn_module(backend="xformers")
+    _force_capability(monkeypatch, (12, 0))
+    m = _fake_model()
+    samplers._funpack_install_mask_safe_attention(m)
+    assert (m.model_options["transformer_options"]["optimized_attention_override"]
+            is samplers._funpack_mask_safe_attention_override)
+
+
+def test_no_install_on_supported_capability(monkeypatch):
+    _install_attn_module(backend="xformers")
+    _force_capability(monkeypatch, (9, 0))  # Hopper — xformers tensor-bias kernels OK
+    m = _fake_model()
+    samplers._funpack_install_mask_safe_attention(m)
+    assert "optimized_attention_override" not in m.model_options.get("transformer_options", {})
+
+
+def test_no_install_when_backend_not_xformers(monkeypatch):
+    _install_attn_module(backend="pytorch")  # user already on SDPA
+    _force_capability(monkeypatch, (12, 0))
+    m = _fake_model()
+    samplers._funpack_install_mask_safe_attention(m)
+    assert "optimized_attention_override" not in m.model_options.get("transformer_options", {})
+
+
+def test_no_install_without_cuda(monkeypatch):
+    _install_attn_module(backend="xformers")
+    _force_capability(monkeypatch, (12, 0), available=False)
+    m = _fake_model()
+    samplers._funpack_install_mask_safe_attention(m)
+    assert "optimized_attention_override" not in m.model_options.get("transformer_options", {})
+
+
+def test_does_not_stomp_existing_override(monkeypatch):
+    _install_attn_module(backend="xformers")
+    _force_capability(monkeypatch, (12, 0))
+    other = lambda func, *a, **k: func(*a, **k)
+    m = types.SimpleNamespace(model_options={"transformer_options": {"optimized_attention_override": other}})
+    samplers._funpack_install_mask_safe_attention(m)
+    assert m.model_options["transformer_options"]["optimized_attention_override"] is other
+
+
+def test_idempotent(monkeypatch):
+    _install_attn_module(backend="xformers")
+    _force_capability(monkeypatch, (12, 0))
+    m = _fake_model()
+    samplers._funpack_install_mask_safe_attention(m)
+    samplers._funpack_install_mask_safe_attention(m)
+    assert (m.model_options["transformer_options"]["optimized_attention_override"]
+            is samplers._funpack_mask_safe_attention_override)