Update FastPLMs files

fastplms/attention/__init__.py

@@ -1,6 +1,7 @@
 """Shared attention backends, masks, and optional optimized kernels."""
 
 from ._core import (
+    LEGACY_CHECKPOINT_ATTENTION_BACKENDS,
     VALID_ATTENTION_BACKENDS,
     AttentionBackend,
     BlockMask,
@@ -11,6 +12,7 @@ from ._core import (
     _kernels_flash_varlen_forward,
     _unpad_input,
     bool_to_additive_mask,
+    canonical_checkpoint_attention_backend,
     clear_flex_attention_caches,
     create_block_mask,
     flex_attention,
@@ -36,6 +38,7 @@ from .interfaces import (
 __all__ = [
     "FASTPLMS_ATTENTION_FUNCTIONS",
     "FASTPLMS_ATTENTION_MASKS",
+    "LEGACY_CHECKPOINT_ATTENTION_BACKENDS",
     "VALID_ATTENTION_BACKENDS",
     "AttentionBackend",
     "BlockMask",
@@ -47,6 +50,7 @@ __all__ = [
     "_kernels_flash_varlen_forward",
     "_unpad_input",
     "bool_to_additive_mask",
+    "canonical_checkpoint_attention_backend",
     "clear_flex_attention_caches",
     "create_block_mask",
     "flex_attention",

fastplms/attention/_core.py

@@ -13,6 +13,7 @@ from collections import OrderedDict
 from collections.abc import Callable
 from enum import Enum
 from threading import RLock
+from types import MappingProxyType
 from einops import rearrange
 from torch.nn import functional as F
 
@@ -623,6 +624,29 @@ class AttentionBackend(str, Enum):  # noqa: UP042
 
 VALID_ATTENTION_BACKENDS = tuple(b.value for b in AttentionBackend)
 
+# Official upstream sources predate the Transformers ``flex_attention`` name and
+# store ``"flex"`` in their checkpoint configurations. The two names select the
+# identical implementation, so translating one to the other is a rename rather
+# than a backend substitution. ``"flash"`` is deliberately absent: it does not
+# identify a FlashAttention version, and guessing one would be a substitution.
+LEGACY_CHECKPOINT_ATTENTION_BACKENDS = MappingProxyType(
+    {"flex": AttentionBackend.FLEX_ATTENTION.value}
+)
+
+
+def canonical_checkpoint_attention_backend(stored_backend: str | None) -> str | None:
+    """Translate a stored checkpoint backend name to its canonical spelling.
+
+    Call this where a serialized configuration enters FastPLMs. Imperative
+    selection through ``set_attn_implementation`` stays strict and keeps
+    rejecting the historical spellings. An unmapped name passes through so the
+    backend resolver owns the rejection.
+    """
+
+    if stored_backend is None:
+        return None
+    return LEGACY_CHECKPOINT_ATTENTION_BACKENDS.get(stored_backend, stored_backend)
+
 
 def warn_attention_backend_fallback(
     requested_backend: str | AttentionBackend,

fastplms/attention/interfaces.py

@@ -10,6 +10,7 @@ from transformers import AttentionInterface, AttentionMaskInterface
 
 from ._core import (
     AttentionBackend,
+    canonical_checkpoint_attention_backend,
     get_attn_implementation,
     kernels_flash_attention_func,
     resolve_attention_backend,
@@ -146,26 +147,32 @@ class FastPLMsAttentionMixin:
     def __init__(self, config, *args: Any, **kwargs: Any) -> None:
         sentinel = object()
         internal = getattr(config, "_attn_implementation_internal", sentinel)
-        canonical = (
+        stored = (
             getattr(config, "_attn_implementation", None) if internal is sentinel else internal
         )
         legacy = getattr(config, "attn_backend", None)
-        requested = canonical if canonical is not None else legacy
+        requested = stored if stored is not None else legacy
         if requested is not None:
             if not isinstance(requested, str):
                 raise TypeError(
                     "The configured attention implementation must be a string or None; "
                     f"received {type(requested).__name__}."
                 )
-            self._validate_attention_name(requested)
+            # A serialized configuration can name a backend with the historical
+            # spelling used by the official source it was converted from. That
+            # names the same implementation, so translate it here rather than
+            # rejecting a checkpoint that asked for an implementation FastPLMs has.
+            canonical = canonical_checkpoint_attention_backend(requested)
+            self._validate_attention_name(canonical)
             # ``PreTrainedModel.__init__`` resolves a missing Transformers
             # implementation to the family default.  Legacy FastPLMs configs
             # persist their explicit choice in ``attn_backend``, so forward it
             # into the canonical Transformers field before the base class can
-            # replace it with SDPA.  A non-None canonical value still wins,
-            # including an explicit ``attn_implementation=...`` load override.
-            if canonical is None and legacy is not None:
-                set_config_attn_implementation(config, legacy)
+            # replace it with SDPA.  A stored canonical value already agrees and
+            # is left untouched, including an explicit
+            # ``attn_implementation=...`` load override.
+            if canonical != stored:
+                set_config_attn_implementation(config, canonical)
         super().__init__(config, *args, **kwargs)
         # Transformers resolves an unspecified implementation during the base
         # model initialization. Synchronize that choice before family layers

fastplms/models.toml

@@ -1,7 +1,7 @@
 schema_version = 1
 legal_files = [
   "LICENSE=sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
-  "THIRD_PARTY_NOTICES.md=sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
+  "THIRD_PARTY_NOTICES.md=sha256:d35e506b728868b52290672d54a89c6aade09529af99dbc8a4db72d9e9ca3460",
 ]
 
 [[attention_kernels]]

fastplms/models/esmfold2/configuration_esmfold2.py

@@ -21,6 +21,9 @@ from typing import Any, TypeVar, cast
 
 from transformers.configuration_utils import PretrainedConfig
 
+from fastplms.attention import canonical_checkpoint_attention_backend
+
+
 _ESMC_ATTENTION_IMPLEMENTATIONS = frozenset({"eager", "flex_attention", "sdpa"})
 _ESMC_PRECISIONS = frozenset({"auto", "bf16", "fp32", "fp8"})
 
@@ -56,7 +59,7 @@ def normalize_esmc_attention_implementation(
                 "ESMFold2 has one ESMC attention backbone; use a string or {'': implementation}."
             )
         implementation = implementation[""]
-    canonical = "flex_attention" if implementation == "flex" else implementation
+    canonical = canonical_checkpoint_attention_backend(implementation)
     if canonical is not None and canonical not in _ESMC_ATTENTION_IMPLEMENTATIONS:
         expected = sorted(_ESMC_ATTENTION_IMPLEMENTATIONS)
         raise ValueError(

modeling_fastplms.py

@@ -13,7 +13,7 @@ from zipfile import ZIP_DEFLATED, ZipFile
 
 from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
 
-if RUNTIME_HASH != "529f5db5434f1737911b7480083002391e7d8a0bc56fec771aa602beb7721475":
+if RUNTIME_HASH != "c76f023da79d5db93093291f4f3858a0e6cde4d9b4fa61ade779cf941f2175da":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []