longcat: pick the attention backend that is installed (v0.2.68)

dit/config.json asks for enable_flashattn2, but requirements-longcat.txt ships
xformers — flash-attn 2.7.4.post1 is a source build needing ninja and a CUDA
toolchain. Upstream's Attention has no fallback: anything unselected reaches
`raise RuntimeError("Unsupported attention operations.")`, and the FA2 branch
imports flash_attn, which is not in the venv.

--attention was only ever written into the environment. Nothing read
LONGCAT_ATTENTION, so the checkpoint's config won every time.

Now both load paths — bf16 through from_pretrained and INT8 through the config
dict — are built with the backend that is actually importable, preferring what
was asked for and saying so when it falls back. None installed is an error
here, with the reason, rather than 48 blocks into the first step of a load that
already cost minutes.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0185zEBiy3KWB37xNGSrxL7w

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.2.67"
+__version__ = "0.2.68"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

tests/test_longcat_placement.py

@@ -165,3 +165,76 @@ def test_no_cuda_is_reported_not_silently_tolerated(service, monkeypatch):
     service._place_pipeline(pipe, "model")
 
     assert pipe.dit.device == "cpu"
+
+
+# ---------------------------------------------------------------- attention backend
+# The checkpoint config asks for FlashAttention-2; the venv installs xformers, because
+# flash-attn is a source build. Upstream's attention raises on anything unselected, and
+# the --attention setting was written to the environment and never read — so the config
+# won and imported a module that is not there.
+
+def _fake_find_spec(installed):
+    def find_spec(name, *a, **k):
+        return object() if name in installed else None
+    return find_spec
+
+
+@pytest.fixture
+def pick(service, monkeypatch):
+    import importlib.util
+
+    def choose(installed, want=None):
+        monkeypatch.setattr(importlib.util, "find_spec", _fake_find_spec(installed))
+        if want is None:
+            monkeypatch.delenv("LONGCAT_ATTENTION", raising=False)
+        else:
+            monkeypatch.setenv("LONGCAT_ATTENTION", want)
+        return service._attention_kwargs()
+    return choose
+
+
+def test_exactly_one_backend_is_ever_enabled(pick):
+    """Two enabled flags would silently take whichever branch is checked first."""
+    flags = pick({"xformers", "flash_attn"})
+    assert sum(bool(v) for v in flags.values()) == 1
+    assert set(flags) == {"enable_flashattn3", "enable_flashattn2", "enable_xformers"}
+
+
+def test_it_picks_what_is_installed_not_what_the_checkpoint_asks_for(pick):
+    """dit/config.json says enable_flashattn2: true and flash_attn is not in the venv."""
+    flags = pick({"xformers"})
+    assert flags["enable_xformers"] is True
+    assert flags["enable_flashattn2"] is False
+
+
+def test_an_explicit_choice_is_honoured_when_it_is_available(pick):
+    flags = pick({"xformers", "flash_attn"}, want="flash")
+    assert flags["enable_flashattn2"] is True
+    assert flags["enable_xformers"] is False
+
+
+def test_asking_for_a_backend_that_is_absent_falls_back(pick):
+    flags = pick({"xformers"}, want="flash")
+    assert flags["enable_xformers"] is True
+
+
+@pytest.mark.parametrize("alias", ["flashattn2", "flash2", "fa2"])
+def test_the_usual_names_for_flashattention_all_work(pick, alias):
+    """A config saying "flashattn2" should not silently land on xformers."""
+    flags = pick({"xformers", "flash_attn"}, want=alias)
+    assert flags["enable_flashattn2"] is True
+
+
+def test_no_backend_at_all_is_an_error_here_not_deep_in_the_first_step(pick):
+    """Upstream's message is "Unsupported attention operations." from inside attention,
+    48 blocks into a load that already cost minutes."""
+    with pytest.raises(RuntimeError, match="no usable attention backend"):
+        pick(set())
+
+
+def test_both_load_paths_get_the_backend(service):
+    """bf16 via from_pretrained and INT8 via the config dict — a fix to one only would
+    leave the other raising on the first attention call."""
+    src = SERVICE.read_text(encoding="utf-8")
+    body = src[src.index("def load_pipeline"):src.index("def _attention_kwargs")]
+    assert body.count("_attention_kwargs()") == 2

tools/longcat_service.py

@@ -126,7 +126,8 @@ def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
                                          cp_split_hw=_state.get("cp_split_hw"))
             else:
                 dit = _load_quantized_base_dit(
-                    root, LC.INT8_SUBDIR, cp_split_hw=_state.get("cp_split_hw"))
+                    root, LC.INT8_SUBDIR, cp_split_hw=_state.get("cp_split_hw"),
+                    **_attention_kwargs())
             log(f"loaded the INT8 DiT from {LC.INT8_SUBDIR}/"
                 + (f" (avatar family '{family}')" if family else " (base)"))
         else:
@@ -144,7 +145,7 @@ def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
                     log("WARNING: this torch has no float8_e4m3fn — loading in "
                         f"{dtype} instead")
             dit = LongCatVideoTransformer3DModel.from_pretrained(
-                root, subfolder="dit", **kw)
+                root, subfolder="dit", **kw, **_attention_kwargs())
 
         if family:
             # The avatar families are a DIFFERENT pipeline class with its own methods
@@ -187,6 +188,53 @@ def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
         return pipe
 
 
+def _attention_kwargs():
+    """Which attention backend the DiT should be built with.
+
+    The checkpoint's own dit/config.json asks for FlashAttention-2, but the venv
+    installs xformers: flash-attn 2.7.4.post1 is a source build needing ninja and a
+    CUDA toolchain, so requirements-longcat.txt ships xformers instead. Upstream's
+    Attention has NO fallback — anything unselected hits
+    `raise RuntimeError("Unsupported attention operations.")` — and the --attention
+    setting was only ever written into the environment, never read. So the config won
+    and the first attention call imported a module that is not installed.
+
+    Pick from what is importable, preferring what was asked for.
+    """
+    import importlib.util
+
+    def have(mod):
+        try:
+            return importlib.util.find_spec(mod) is not None
+        except (ImportError, ValueError):
+            return False
+
+    backends = {
+        "flash3": ("flash_attn_interface", {"enable_flashattn3": True}),
+        "flash": ("flash_attn", {"enable_flashattn2": True}),
+        "xformers": ("xformers", {"enable_xformers": True}),
+    }
+    want = (os.environ.get("LONGCAT_ATTENTION") or "xformers").strip().lower()
+    want = {"flashattn2": "flash", "flash2": "flash", "fa2": "flash",
+            "flashattn3": "flash3", "fa3": "flash3"}.get(want, want)
+    order = [want] + [k for k in ("xformers", "flash", "flash3") if k != want]
+
+    off = {"enable_flashattn3": False, "enable_flashattn2": False,
+           "enable_xformers": False}
+    for name in order:
+        entry = backends.get(name)
+        if entry and have(entry[0]):
+            if name != want:
+                log(f"WARNING: attention '{want}' is not installed in this venv; "
+                    f"using '{name}'")
+            log(f"attention backend: {name}")
+            return dict(off, **entry[1])
+    raise RuntimeError(
+        "no usable attention backend: none of xformers, flash_attn or "
+        "flash_attn_interface is installed in the LongCat venv, and upstream's "
+        "attention has no fallback path")
+
+
 def _place_pipeline(pipe, offload):
     """Put the pipeline on the card, honouring the offload mode.