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>
2026-06-28 21:41:43 +03:00
parent 4efe3ba851
commit ff224cbffb
2 changed files with 200 additions and 0 deletions
@@ -600,6 +600,62 @@ def _apply_quality_sharpness(denoised, prev_denoised, sharpness):
        return denoised


+# ---------------------------------------------------------------------------
+# Blackwell (sm_120) xformers masked-attention fallback
+# ---------------------------------------------------------------------------
+# On GPUs newer than capability (9, 0), xformers' memory_efficient_attention has no
+# kernel that accepts a tensor attn_bias — every backend is rejected ("No operator
+# found ... too new ... attn_bias type is <class 'torch.Tensor'>"). Unmasked attention
+# still works, which is why i2v ANCHOR scenes generate fine but GUIDE scenes crash: the
+# LTX guide path (comfy/ldm/lightricks/model.py _attention_with_guide_mask) passes a
+# tensor mask. PyTorch SDPA handles masks on every GPU, so when the active backend is
+# xformers on such a device we route only the MASKED calls to attention_pytorch via
+# ComfyUI's per-call optimized_attention_override hook (wrap_attn). Unmasked calls stay
+# on the fast xformers path. This is the programmatic equivalent of the user's
+# --use-pytorch-cross-attention, scoped to masked attention and installed ONLY when this
+# exact failure condition is detected, so other setups are untouched.
+def _funpack_mask_safe_attention_override(func, *args, **kwargs):
+    # `func` is the raw backend ComfyUI selected (attention_xformers). wrap_attn has
+    # already stamped _inside_attn_wrapper into kwargs before calling us, so invoking
+    # attention_pytorch here cannot re-enter the override (no recursion).
+    mask = kwargs.get("mask")
+    if mask is None and len(args) >= 5:
+        mask = args[4]  # (q, k, v, heads, mask, ...)
+    if mask is not None:
+        import comfy.ldm.modules.attention as _am
+        return _am.attention_pytorch(*args, **kwargs)
+    return func(*args, **kwargs)
+
+
+def _funpack_install_mask_safe_attention(model):
+    """When the active attention backend is xformers on a GPU too new for its tensor-bias
+    kernels (capability > (9, 0), e.g. Blackwell sm_120), install a per-call override so
+    MASKED attention (the LTX guide path) falls back to SDPA. No-op on supported GPUs, when
+    ComfyUI already uses a mask-capable backend, or when another override is present.
+    Idempotent; never raises."""
+    try:
+        import comfy.ldm.modules.attention as _am
+        if getattr(_am, "optimized_attention", None) is not getattr(_am, "attention_xformers", None):
+            return  # SDPA / sage / flash / split selected — masks already handled
+        if not torch.cuda.is_available():
+            return
+        major, minor = torch.cuda.get_device_capability()
+        if major <= 9:
+            return  # xformers cutlass/flash tensor-bias kernels support cap <= (9, 0)
+        to = model.model_options.setdefault("transformer_options", {})
+        existing = to.get("optimized_attention_override")
+        if existing is _funpack_mask_safe_attention_override:
+            return  # already installed (idempotent across scenes/runs)
+        if existing is not None:
+            return  # respect another tool's override; don't stomp it
+        to["optimized_attention_override"] = _funpack_mask_safe_attention_override
+        print(f"[FunPack AV] xformers active on capability {major}.{minor} GPU (too new for "
+              "its tensor-bias kernels): routing MASKED attention to PyTorch SDPA so guide "
+              "scenes don't crash. Unmasked attention stays on xformers.")
+    except Exception as _e:
+        print(f"[FunPack AV] mask-safe attention install skipped: {_e}")
+
+
 # ---------------------------------------------------------------------------
 # KV-Lock: variance-gated scheduler for BachVid K/V injection (arxiv 2603.09657)
 #
@@ -3768,6 +3824,12 @@ class FunPackLTXAVSceneChainSampler:
        except Exception as _e:
            print(f"[FunPackLTXAVSceneChainSampler] hook strip failed: {_e}")

+        # Blackwell (sm_120) GPUs can't run xformers attention with a tensor mask; the LTX
+        # guide path uses one, so anchor scenes generate but guide scenes crash. Route masked
+        # attention to SDPA when that exact combo is detected (no-op otherwise). Threaded via
+        # transformer_options so it reaches every scene's model forward.
+        _funpack_install_mask_safe_attention(model)
+
        # Decode-time noise (folded in from LTXV's Set VAE Decoder Noise per the boundary law:
        # the Chain Sampler owns IMAGES decode, so this lives here, not on a separate node).
        # Stamp settings onto a shallow copy of the VAE so we never mutate the shared input.
@@ -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)