Update FastPLMs runtime and cards from 328deef (files only)

README.md

@@ -201,6 +201,39 @@ not target coordinates.
 Predicted coordinates and confidence scores are outputs and do not establish
 biochemical activity.
 
+## Folding speed settings
+
+Two runtime settings trade memory or exactness for speed on long proteins. They
+need no extra package and no compilation, and neither is stored in the
+configuration.
+
+```python
+model.set_chunk_size(None)            # unchunked pair updates
+model.set_atom_attention("windowed")  # the official flash-attn atom window, through PyTorch
+```
+
+`set_chunk_size(None)` removes the row chunking of the pair-update blocks, which
+costs most of a long fold's time on a data-center GPU and saves little peak
+memory; pass a chunk such as 512 when the unchunked fold does not fit.
+`set_atom_attention("windowed")` restricts each atom to 64 real neighbors on
+each side, as the official model does when flash-attn is installed. It needs
+CUDA and changes numerical output, within sampling spread on the measured
+panel. The
+[ESMFold2 guide](https://github.com/Synthyra/FastPLMs/blob/main/docs/esmfold2.md#measured-folding-cost)
+records the conditions, the dense-versus-windowed comparison, and the figure.
+
+| Residues | FastPLMs defaults (s) | Optimized (s) | Official implementation (s) |
+| ---: | ---: | ---: | ---: |
+| 256 | 1.8 | 1.6 | 1.9 |
+| 1,024 | 81 | 23 | 87 |
+| 2,048 | 596 | 139 with 512-row chunks | out of memory |
+
+Measured on one NVIDIA H100 80GB HBM3 with PyTorch 2.13.0+cu130: one
+fixed pseudo-random protein per length, 3 trunk loops,
+50 requested sampling steps under the official noise cap,
+1 diffusion sample, BF16 autocast over FP32 folding
+parameters, median of end-to-end folds. "Defaults" changes no setting.
+
 ## Learned representation and ESMC precision
 
 ESMFold2 applies its learned state mixture and projection as

fastplms/attention/__init__.py

@@ -1,10 +1,17 @@
 """Shared attention backends, masks, and optional optimized kernels."""
 
+from ._auto import (
+    AUTO_ATTENTION,
+    AttentionCandidate,
+    AttentionExecutionContext,
+    AttentionResolution,
+)
 from ._core import (
     LEGACY_CHECKPOINT_ATTENTION_BACKENDS,
     VALID_ATTENTION_BACKENDS,
     AttentionBackend,
     BlockMask,
+    FlashPaddingLayout,
     _ensure_flash_kernels_loaded,
     _get_flex_attention_fn,
     _get_flex_block_mask,
@@ -18,6 +25,7 @@ from ._core import (
     flex_attention,
     get_attention_mask,
     get_attn_implementation,
+    get_flash_padding_layout,
     index_first_axis,
     index_put_first_axis,
     kernels_flash_attention_func,
@@ -36,13 +44,18 @@ from .interfaces import (
 
 
 __all__ = [
+    "AUTO_ATTENTION",
     "FASTPLMS_ATTENTION_FUNCTIONS",
     "FASTPLMS_ATTENTION_MASKS",
     "LEGACY_CHECKPOINT_ATTENTION_BACKENDS",
     "VALID_ATTENTION_BACKENDS",
     "AttentionBackend",
+    "AttentionCandidate",
+    "AttentionExecutionContext",
+    "AttentionResolution",
     "BlockMask",
     "FastPLMsAttentionMixin",
+    "FlashPaddingLayout",
     "_ensure_flash_kernels_loaded",
     "_get_flex_attention_fn",
     "_get_flex_block_mask",
@@ -56,6 +69,7 @@ __all__ = [
     "flex_attention",
     "get_attention_mask",
     "get_attn_implementation",
+    "get_flash_padding_layout",
     "index_first_axis",
     "index_put_first_axis",
     "kernels_flash_attention_func",

fastplms/attention/_auto.py

@@ -0,0 +1,182 @@
+"""Opt-in automatic selection of an attention implementation.
+
+``attn_implementation="auto"`` is a request, not a backend. Each family lists
+its implementations in a preference order backed by measured evidence, and
+FastPLMs configures the first one that this machine can execute. After that the
+model holds a named implementation. Configuration files and embedding
+fingerprints record that name and never the word ``auto``.
+
+Eager attention, SDPA, and Flex attention can be judged when the model is built.
+FlashAttention depends on the device and on the dtype that Q, K, and V will
+have, so a preference order that contains it is resolved at the first forward or
+by an explicit ``resolve_attn_implementation`` call.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from dataclasses import dataclass
+
+from ._core import _ensure_flash_kernels_loaded, resolve_attention_backend
+from ._kernel_lock import require_kernels_package
+
+
+AUTO_ATTENTION = "auto"
+_FLASH_IMPLEMENTATIONS = frozenset({"flash_attention_2", "flash_attention_3"})
+_DTYPE_NAMES = {
+    torch.float32: "float32",
+    torch.bfloat16: "bfloat16",
+    torch.float16: "float16",
+}
+
+
+@dataclass(frozen=True)
+class AttentionExecutionContext:
+    """The device and dtype that the attention inputs of a forward will have."""
+
+    device: torch.device
+    dtype: torch.dtype
+
+
+@dataclass(frozen=True)
+class AttentionCandidate:
+    """One implementation from the preference order and why it was or was not usable."""
+
+    implementation: str
+    usable: bool
+    reason: str
+
+
+@dataclass(frozen=True)
+class AttentionResolution:
+    """The outcome of one ``auto`` request.
+
+    While ``deferred`` is true the model runs ``resolved`` provisionally, and the
+    first forward or ``resolve_attn_implementation`` replaces this record.
+    """
+
+    requested: str
+    resolved: str
+    candidates: tuple[AttentionCandidate, ...]
+    context: AttentionExecutionContext | None
+    deferred: bool
+
+
+def needs_execution_context(order: tuple[str, ...]) -> bool:
+    return any(implementation in _FLASH_IMPLEMENTATIONS for implementation in order)
+
+
+def provisional_implementation(order: tuple[str, ...]) -> str:
+    """Return the implementation a model runs until its context is known."""
+    for implementation in order:
+        if implementation not in _FLASH_IMPLEMENTATIONS:
+            return implementation
+    raise ValueError(f"The automatic attention order {order} has no context-free implementation.")
+
+
+def attention_execution_context(
+    module: torch.nn.Module,
+    device: torch.device | str | None = None,
+    dtype: torch.dtype | None = None,
+) -> AttentionExecutionContext:
+    """Describe where the next forward of ``module`` will run its attention.
+
+    Under CUDA autocast FP32 parameters produce autocast-dtype Q, K, and V, so
+    the autocast dtype is the one that decides kernel eligibility.
+    """
+    parameter = next(module.parameters(), None)
+    if parameter is None:
+        raise RuntimeError("Automatic attention selection requires a model with parameters.")
+    resolved_device = parameter.device if device is None else torch.device(device)
+    if dtype is not None:
+        return AttentionExecutionContext(resolved_device, dtype)
+    if resolved_device.type == "cuda" and torch.is_autocast_enabled("cuda"):
+        return AttentionExecutionContext(resolved_device, torch.get_autocast_dtype("cuda"))
+    return AttentionExecutionContext(resolved_device, parameter.dtype)
+
+
+def _unusable(implementation: str, reason: str) -> AttentionCandidate:
+    return AttentionCandidate(implementation, False, reason)
+
+
+def _flash_candidate(implementation: str, context: AttentionExecutionContext) -> AttentionCandidate:
+    """Judge a FlashAttention kernel, leaving the possible download for the last gate."""
+    from fastplms.registry import get_model_registry
+
+    kernel_spec = get_model_registry().attention_kernels[implementation]
+    if context.device.type != "cuda":
+        return _unusable(
+            implementation, f"It requires a CUDA device; the model is on {context.device}."
+        )
+    dtype_name = _DTYPE_NAMES.get(context.dtype, str(context.dtype))
+    if dtype_name not in kernel_spec.dtypes:
+        supported = ", ".join(kernel_spec.dtypes)
+        return _unusable(
+            implementation,
+            f"It supports only {supported}; the attention inputs would be {dtype_name}. "
+            "Use CUDA BF16 autocast or BF16 weights.",
+        )
+    capability = torch.cuda.get_device_capability(context.device)
+    if capability < kernel_spec.min_cuda_capability:
+        required = ".".join(str(part) for part in kernel_spec.min_cuda_capability)
+        observed = ".".join(str(part) for part in capability)
+        return _unusable(
+            implementation,
+            f"It requires CUDA compute capability {required} or newer; this GPU has {observed}.",
+        )
+    try:
+        require_kernels_package()
+        _ensure_flash_kernels_loaded(implementation)
+    except RuntimeError as error:
+        return _unusable(implementation, str(error))
+    return AttentionCandidate(implementation, True, "The manifest-locked kernel loaded.")
+
+
+def _candidate(
+    implementation: str, context: AttentionExecutionContext | None
+) -> AttentionCandidate:
+    if implementation in _FLASH_IMPLEMENTATIONS:
+        if context is None:
+            return _unusable(implementation, "It needs a device and dtype, which are not known.")
+        return _flash_candidate(implementation, context)
+    try:
+        # Flex attention is the one device-independent backend a PyTorch build can lack.
+        resolve_attention_backend(implementation)
+    except RuntimeError as error:
+        return _unusable(implementation, str(error))
+    return AttentionCandidate(implementation, True, "It runs on every supported device and dtype.")
+
+
+def resolve_auto_attention(
+    order: tuple[str, ...],
+    context: AttentionExecutionContext | None,
+) -> AttentionResolution:
+    """Select the first implementation in ``order`` that can execute in ``context``."""
+    candidates: list[AttentionCandidate] = []
+    for implementation in order:
+        candidate = _candidate(implementation, context)
+        candidates.append(candidate)
+        if candidate.usable:
+            return AttentionResolution(
+                requested=AUTO_ATTENTION,
+                resolved=implementation,
+                candidates=tuple(candidates),
+                context=context,
+                deferred=False,
+            )
+    reasons = "; ".join(
+        f"{candidate.implementation}: {candidate.reason}" for candidate in candidates
+    )
+    raise RuntimeError(f"No implementation in the automatic attention order is usable. {reasons}")
+
+
+def deferred_resolution(order: tuple[str, ...]) -> AttentionResolution:
+    """Record an ``auto`` request that waits for its execution context."""
+    return AttentionResolution(
+        requested=AUTO_ATTENTION,
+        resolved=provisional_implementation(order),
+        candidates=(),
+        context=None,
+        deferred=True,
+    )

fastplms/attention/_core.py

@@ -11,6 +11,7 @@ import warnings
 import torch
 from collections import OrderedDict
 from collections.abc import Callable
+from dataclasses import dataclass
 from enum import Enum
 from threading import RLock
 from types import MappingProxyType
@@ -440,12 +441,53 @@ index_first_axis = IndexFirstAxis.apply
 index_put_first_axis = IndexPutFirstAxis.apply
 
 
+def _select_first_axis(states: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
+    """Copy the rows at ``indices``, through the autograd wrapper only when it is needed.
+
+    Without a gradient to route, the wrapper and its reshapes are host overhead
+    that every layer pays, and plain indexing copies the same values.
+    """
+    # states: (n, ...); indices: (m,)
+    if states.requires_grad:
+        selected: torch.Tensor = index_first_axis(states, indices)  # (m, ...)
+        return selected
+    return states[indices]  # (m, ...)
+
+
 def pad_input(
     hidden_states: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int
 ) -> torch.Tensor:
     # hidden_states: (t, ...); indices: (t,)
-    output = index_put_first_axis(hidden_states, indices, batch * seqlen)  # (b * l, ...)
-    return rearrange(output, "(b s) ... -> b s ...", b=batch)  # (b, l, ...)
+    if hidden_states.requires_grad:
+        output = index_put_first_axis(hidden_states, indices, batch * seqlen)  # (b * l, ...)
+        return rearrange(output, "(b s) ... -> b s ...", b=batch)  # (b, l, ...)
+    output = hidden_states.new_zeros(batch * seqlen, *hidden_states.shape[1:])  # (b * l, ...)
+    output[indices] = hidden_states
+    return output.view(batch, seqlen, *hidden_states.shape[1:])  # (b, l, ...)
+
+
+@dataclass(frozen=True)
+class FlashPaddingLayout:
+    """Varlen metadata shared by every layer of one padded self-attention forward.
+
+    The token indices and the longest row each cost a host synchronization to
+    derive. An encoder builds this record once per forward so its layers do not
+    repeat that work. It is valid only for the mask it was built from.
+    """
+
+    attention_mask_2d: torch.Tensor  # (b, l), bool; True marks a real token
+    indices: torch.Tensor  # (t,), positions of the t real tokens in the flat (b * l) axis
+    cu_seqlens: torch.Tensor  # (b + 1,), int32 cumulative row lengths
+    max_seqlen: int  # longest row
+
+
+def _flash_padding_layout(attention_mask_2d: torch.Tensor) -> FlashPaddingLayout:
+    # attention_mask_2d: (b, l), bool
+    seqlens = attention_mask_2d.sum(dim=1).int()  # (b,)
+    cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0))  # (b + 1,)
+    max_seqlen = int(seqlens.max().item())
+    indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten()  # (t,)
+    return FlashPaddingLayout(attention_mask_2d, indices, cu_seqlens, max_seqlen)
 
 
 def _unpad_input(
@@ -453,6 +495,7 @@ def _unpad_input(
     key_layer: torch.Tensor,
     value_layer: torch.Tensor,
     attention_mask_2d: torch.Tensor,
+    padding_layout: FlashPaddingLayout | None = None,
 ) -> tuple[
     torch.Tensor,
     torch.Tensor,
@@ -463,17 +506,18 @@ def _unpad_input(
 ]:
     # query_layer, key_layer, value_layer: (b, l, h, d); attention_mask_2d: (b, l)
     batch_size, seq_len, num_heads, head_dim = query_layer.shape
-    seqlens = attention_mask_2d.sum(dim=1).int()  # (b,)
-    cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0))  # (b + 1,)
-    max_seqlen = int(seqlens.max().item())
-    indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten()  # (t,)
-    query_layer = index_first_axis(  # (t, h, d)
+    if padding_layout is None:
+        padding_layout = _flash_padding_layout(attention_mask_2d)
+    indices = padding_layout.indices  # (t,)
+    cu_seqlens = padding_layout.cu_seqlens  # (b + 1,)
+    max_seqlen = padding_layout.max_seqlen
+    query_layer = _select_first_axis(  # (t, h, d)
         query_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices
     )
-    key_layer = index_first_axis(  # (t, h, d)
+    key_layer = _select_first_axis(  # (t, h, d)
         key_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices
     )
-    value_layer = index_first_axis(  # (t, h, d)
+    value_layer = _select_first_axis(  # (t, h, d)
         value_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices
     )
     return (
@@ -513,6 +557,25 @@ def _validate_flash_padding_mask(
     return attention_mask_2d.to(dtype=torch.bool)  # (b, l)
 
 
+def _validate_flash_padding_layout(
+    padding_layout: FlashPaddingLayout,
+    attention_mask_2d: torch.Tensor | None,
+) -> None:
+    """Reject a layout that cannot belong to this call without reading tensor values."""
+
+    # attention_mask_2d: (b, l) or None
+    if attention_mask_2d is None:
+        raise ValueError("A FlashAttention padding layout requires the mask it was built from.")
+    layout_mask = padding_layout.attention_mask_2d  # (b, l)
+    if layout_mask.shape != attention_mask_2d.shape:
+        raise ValueError(
+            "FlashAttention padding layout was built for mask shape "
+            f"{tuple(layout_mask.shape)}, but this call uses {tuple(attention_mask_2d.shape)}."
+        )
+    if layout_mask.device != attention_mask_2d.device:
+        raise ValueError("FlashAttention padding layout and padding mask must share a device.")
+
+
 def kernels_flash_attention_func(
     query_states: torch.Tensor,
     key_states: torch.Tensor,
@@ -521,6 +584,7 @@ def kernels_flash_attention_func(
     causal: bool = False,
     softmax_scale: float | None = None,
     implementation: str = "flash_attention_3",
+    padding_layout: FlashPaddingLayout | None = None,
 ) -> torch.Tensor:
     """Public flash-attention entry point with optional padding handling.
 
@@ -533,6 +597,10 @@ def kernels_flash_attention_func(
     before calling this function (ESM2, DPLM, DPLM2, E1, and ESMFold do), pass
     `softmax_scale=1.0`. Otherwise the flash kernel applies its default scale
     again, yielding an effective `1/head_dim` scale that drifts across layers.
+
+    `padding_layout` is the record `get_flash_padding_layout` built from
+    `attention_mask_2d`. Pass both so that a multi-layer encoder derives the
+    varlen metadata once per forward. The output is identical without it.
     """
     # query_states, key_states, value_states: (b, l, h, d)
     # attention_mask_2d: (b, l) or None
@@ -559,6 +627,8 @@ def kernels_flash_attention_func(
             value_states,
             attention_mask_2d,
         )
+    if padding_layout is not None:
+        _validate_flash_padding_layout(padding_layout, attention_mask_2d)
     _ensure_flash_kernels_loaded(implementation)
     if attention_mask_2d is not None:
         batch_size, q_len = query_states.shape[:2]
@@ -574,6 +644,7 @@ def kernels_flash_attention_func(
             key_states,
             value_states,
             attention_mask_2d,
+            padding_layout,
         )
         attn_output_unpad = _kernels_flash_varlen_forward(  # (t, h, d)
             query_states=query_states,
@@ -587,8 +658,8 @@ def kernels_flash_attention_func(
             softmax_scale=softmax_scale,
             implementation=implementation,
         )
-        output = pad_input(attn_output_unpad, indices_q, batch_size, q_len)  # (b, l, h, d)
-        return output.masked_fill(~attention_mask_2d[:, :, None, None], 0)  # (b, l, h, d)
+        # pad_input scatters the real tokens into zeros, so padded positions are already zero.
+        return pad_input(attn_output_unpad, indices_q, batch_size, q_len)  # (b, l, h, d)
     else:
         return _kernels_flash_forward(  # (b, l, h, d)
             query_states=query_states,
@@ -803,6 +874,21 @@ def get_attention_mask(
     return attention_mask_2d, attention_mask_4d, None  # (b, l), (b, 1, 1, l), None
 
 
+def get_flash_padding_layout(
+    effective_backend: AttentionBackend,
+    attention_mask_2d: torch.Tensor | None,
+) -> FlashPaddingLayout | None:
+    """Build the FlashAttention varlen metadata once for all encoder layers.
+
+    `attention_mask_2d` is the bool mask that `get_attention_mask` returned.
+    Returns None when the call does not run a padded FlashAttention kernel.
+    """
+    # attention_mask_2d: (b, l) or None
+    if attention_mask_2d is None or not resolve_attention_backend(effective_backend).is_flash:
+        return None
+    return _flash_padding_layout(attention_mask_2d)
+
+
 def bool_to_additive_mask(
     bool_mask: torch.Tensor,
     dtype: torch.dtype,

fastplms/attention/interfaces.py

@@ -8,6 +8,15 @@ from functools import partial
 from typing import Any
 from transformers import AttentionInterface, AttentionMaskInterface
 
+from ._auto import (
+    AUTO_ATTENTION,
+    AttentionResolution,
+    attention_execution_context,
+    deferred_resolution,
+    needs_execution_context,
+    provisional_implementation,
+    resolve_auto_attention,
+)
 from ._core import (
     AttentionBackend,
     canonical_checkpoint_attention_backend,
@@ -95,6 +104,13 @@ class FastPLMsAttentionMixin:
         "sdpa",
         "flex_attention",
     )
+    # Preference order for ``attn_implementation="auto"``, mirrored from the
+    # family's ``attention_auto_order`` in models.toml. Empty rejects the request.
+    _fastplms_attention_auto_order: tuple[str, ...] = ()
+    # Supplied by the Transformers ``PreTrainedModel`` that follows this mixin in every
+    # MRO. Each family's configuration class adds the ``attn_backend`` field this mixin
+    # reads, and they share no base that declares it, so the boundary is dynamic.
+    config: Any
 
     def _validate_attention_name(self, implementation: str) -> None:
         if implementation not in self._fastplms_attention_implementations:
@@ -147,11 +163,18 @@ class FastPLMsAttentionMixin:
     def __init__(self, config, *args: Any, **kwargs: Any) -> None:
         sentinel = object()
         internal = getattr(config, "_attn_implementation_internal", sentinel)
-        stored = (
-            getattr(config, "_attn_implementation", None) if internal is sentinel else internal
-        )
+        stored = getattr(config, "_attn_implementation", None) if internal is sentinel else internal
         legacy = getattr(config, "attn_backend", None)
         requested = stored if stored is not None else legacy
+        auto_requested = requested == AUTO_ATTENTION
+        serialized_backend: str | None = None
+        if auto_requested:
+            # The configuration never holds ``auto``. Family layers are built on the
+            # provisional implementation, and a saved copy keeps the backend that a
+            # named load of the same checkpoint would have stored.
+            requested = provisional_implementation(self._require_attention_auto_order())
+            serialized_backend = legacy if legacy not in (None, AUTO_ATTENTION) else requested
+            stored = None
         if requested is not None:
             if not isinstance(requested, str):
                 raise TypeError(
@@ -180,6 +203,93 @@ class FastPLMsAttentionMixin:
         resolved = get_attn_implementation(config)
         self._validate_attention_name(resolved)
         set_config_attn_implementation(config, resolved)
+        if auto_requested:
+            self.__dict__["_fastplms_serialized_attn_backend"] = serialized_backend
+            self._begin_auto_attention()
+
+    def _require_attention_auto_order(self) -> tuple[str, ...]:
+        order = self._fastplms_attention_auto_order
+        if not order:
+            raise ValueError(
+                f"{type(self).__name__} does not support attn_implementation='auto'; "
+                f"request one of {self._fastplms_attention_implementations}."
+            )
+        return order
+
+    @property
+    def attention_resolution(self) -> AttentionResolution | None:
+        """The record of an ``auto`` request, or None when a backend was named."""
+        return self.__dict__.get("_fastplms_attention_resolution")
+
+    def _begin_auto_attention(self) -> None:
+        """Resolve now when no candidate needs a device, else at the first forward."""
+        order = self._require_attention_auto_order()
+        self._cancel_pending_auto_attention()
+        if not needs_execution_context(order):
+            resolution = resolve_auto_attention(order, None)
+            self._apply_attn_implementation(resolution.resolved)
+            self.__dict__["_fastplms_attention_resolution"] = resolution
+            return
+        resolution = deferred_resolution(order)
+        self._apply_attn_implementation(resolution.resolved)
+        self.__dict__["_fastplms_attention_resolution"] = resolution
+        # The first forward runs inside the caller's autocast context, which is
+        # what decides FlashAttention eligibility for FP32 parameters.
+        self.__dict__["_fastplms_auto_attention_hook"] = (
+            self._as_module().register_forward_pre_hook(_resolve_auto_attention_before_forward)
+        )
+
+    def _as_module(self) -> torch.nn.Module:
+        if not isinstance(self, torch.nn.Module):
+            raise TypeError(
+                f"{type(self).__name__} must be a torch.nn.Module to defer attention selection."
+            )
+        return self
+
+    def _cancel_pending_auto_attention(self) -> None:
+        hook = self.__dict__.pop("_fastplms_auto_attention_hook", None)
+        if hook is not None:
+            hook.remove()
+
+    def resolve_attn_implementation(
+        self,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> AttentionResolution:
+        """Settle a pending ``auto`` request for the device and dtype of the next forward.
+
+        The first forward does this by itself. Call it earlier, for example before
+        ``torch.compile`` or before fingerprinting an embedding run, and pass
+        ``dtype`` when the forward will run under an autocast context that is not
+        active yet. A settled request returns its record unchanged.
+        """
+        resolution = self.attention_resolution
+        if resolution is None:
+            raise RuntimeError(
+                f"{type(self).__name__} was not configured with attn_implementation='auto'."
+            )
+        if not resolution.deferred:
+            return resolution
+        self._cancel_pending_auto_attention()
+        resolution = resolve_auto_attention(
+            self._require_attention_auto_order(),
+            attention_execution_context(self._as_module(), device=device, dtype=dtype),
+        )
+        self._apply_attn_implementation(resolution.resolved)
+        self.__dict__["_fastplms_attention_resolution"] = resolution
+        return resolution
+
+    def save_pretrained(self, *args: Any, **kwargs: Any) -> Any:
+        """Save without the machine-specific outcome of an ``auto`` request."""
+        # ``save_pretrained`` comes from the ``PreTrainedModel`` later in the MRO.
+        if self.attention_resolution is None:
+            return super().save_pretrained(*args, **kwargs)  # type: ignore[misc]
+        selected_backend = self.config.attn_backend
+        self.config.attn_backend = self.__dict__["_fastplms_serialized_attn_backend"]
+        try:
+            return super().save_pretrained(*args, **kwargs)  # type: ignore[misc]
+        finally:
+            self.config.attn_backend = selected_backend
 
     def set_attn_implementation(
         self,
@@ -194,6 +304,23 @@ class FastPLMsAttentionMixin:
                 raise ValueError(
                     "FastPLMs models have one attention backbone; pass a string or {'': name}."
                 )
+        if attn_implementation == AUTO_ATTENTION:
+            if allow_all_kernels:
+                raise ValueError("FastPLMs does not load external attention kernels.")
+            self.__dict__.setdefault(
+                "_fastplms_serialized_attn_backend", getattr(self.config, "attn_backend", None)
+            )
+            self._begin_auto_attention()
+            return
+        # A named request replaces any earlier automatic selection.
+        self._cancel_pending_auto_attention()
+        self.__dict__.pop("_fastplms_attention_resolution", None)
+        self.__dict__.pop("_fastplms_serialized_attn_backend", None)
+        self._apply_attn_implementation(attn_implementation, allow_all_kernels)
+
+    def _apply_attn_implementation(
+        self, attn_implementation: str, allow_all_kernels: bool = False
+    ) -> None:
         resolved_name = self._check_and_adjust_attn_implementation(
             attn_implementation,
             is_init_check=False,
@@ -213,6 +340,17 @@ class FastPLMsAttentionMixin:
                 )
 
 
+# Selection reads the manifest, can load a kernel, and rewrites layer attributes.
+# It runs eagerly so that a compiled model never traces it.
+@torch.compiler.disable  # type: ignore[untyped-decorator]
+def _resolve_auto_attention_before_forward(
+    module: torch.nn.Module, _arguments: tuple[Any, ...]
+) -> None:
+    if not isinstance(module, FastPLMsAttentionMixin):
+        raise TypeError("The automatic attention hook belongs on a FastPLMs model.")
+    module.resolve_attn_implementation()
+
+
 def validate_transformers_attention_interfaces() -> None:
     """Verify that Transformers exposes functions and masks for every backend.
 

fastplms/embeddings/batches.py

@@ -0,0 +1,378 @@
+"""Execute model-specific batches and return ordered residue-aware CPU tensors."""
+
+from __future__ import annotations
+
+import torch
+from collections.abc import Callable, Iterator, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any
+from torch import Tensor
+
+from .identity import _model_device
+from .inputs import _planned_batches
+from .pooling import Pooler
+from .types import EmbeddingBatch, EmbeddingInput, EmbeddingRecord
+
+
+_MAX_PARTI_RESIDUES = 2_048
+
+
+def _validate_parti_length(M: Tensor) -> None:
+    """Reject an oversized attention graph before model inference."""
+
+    # M: (b, l)
+    n_residues = int(M.to(dtype=torch.int64).sum(dim=1).max().item())
+    if n_residues > _MAX_PARTI_RESIDUES:
+        raise ValueError(f"parti supports at most {_MAX_PARTI_RESIDUES:,} biological residues.")
+
+
+def select_hidden_state_embeddings(
+    last_hidden_state: Tensor,
+    hidden_states: tuple[Tensor, ...] | None,
+    *,
+    hidden_state_index: int = -1,
+    store_all_hidden_states: bool = False,
+) -> Tensor:
+    """Select one hidden state or stack every state without changing values."""
+    # last_hidden_state and each hidden_states entry: (b, l, d)
+    if store_all_hidden_states:
+        if not hidden_states:
+            raise ValueError("store_all_hidden_states requires model hidden states.")
+        # H has shape (b, n, l, d), where n follows the model's output order.
+        return torch.stack(hidden_states, dim=1)  # (b, n, l, d)
+    if hidden_state_index == -1:
+        return last_hidden_state  # (b, l, d)
+    if not hidden_states:
+        raise ValueError("hidden_state_index requires model hidden states.")
+    return hidden_states[hidden_state_index]  # (b, l, d)
+
+
+def _residue_embeddings(X: Tensor, M: Tensor) -> list[Tensor]:
+    """Copy every sample's biological residues to the host in one transfer.
+
+    Boolean indexing packs the selected rows in batch order, so splitting the
+    packed rows by residue count gives the values that indexing each sample
+    would. Each returned tensor owns its storage, as a per-sample copy does.
+    """
+    # X: (b, l, d); M: (b, l)
+    residue_counts = M.sum(dim=1).tolist()  # b counts r_i
+    packed = X[M].detach().cpu()  # (sum of r_i, d)
+    return [sample.clone() for sample in torch.split(packed, residue_counts)]  # each: (r_i, d)
+
+
+@contextmanager
+def _temporary_eval(model: Any) -> Iterator[None]:
+    was_training = getattr(model, "training", None)
+    eval_method = getattr(model, "eval", None)
+    train_method = getattr(model, "train", None)
+    if (
+        not isinstance(was_training, bool)
+        or not callable(eval_method)
+        or not callable(train_method)
+    ):
+        yield
+        return
+    eval_method()
+    try:
+        yield
+    finally:
+        train_method(was_training)
+
+
+def _biological_residue_mask(
+    input_ids: Tensor,
+    attention_mask: Tensor,
+    tokenizer: Any,
+) -> Tensor:
+    """Remove padding and tokenizer-declared special tokens from M."""
+
+    # input_ids, attention_mask: (b, l)
+    M = attention_mask.to(dtype=torch.bool)  # (b, l)
+    special_ids = tuple(int(token_id) for token_id in getattr(tokenizer, "all_special_ids", ()))
+    if special_ids:
+        specials = torch.tensor(  # (n_special,)
+            special_ids,
+            device=input_ids.device,
+            dtype=input_ids.dtype,
+        )
+        M = M & ~torch.isin(input_ids, specials)  # (b, l)
+    return M  # (b, l)
+
+
+def _generic_embedding_batch(
+    model: Any,
+    sequences: list[str],
+    *,
+    tokenizer: Any | None,
+    max_length: int | None,
+    truncate: bool,
+    need_attentions: bool,
+    model_kwargs: dict[str, Any],
+) -> EmbeddingBatch:
+    config = getattr(model, "config", None)
+    model_type = str(getattr(config, "model_type", "")).lower()
+    if tokenizer is None:
+        tokenizer = getattr(model, "tokenizer", None)
+
+    if tokenizer is None and model_type == "e1":
+        output = model._embed(sequences, return_attention_mask=True, **model_kwargs)
+        if not isinstance(output, tuple) or len(output) != 2:
+            raise TypeError("E1 _embed must return (X, residue_mask).")
+        X, M = output  # (b, l, d), (b, l)
+        preparer = getattr(model, "prep_tokens", None)
+        if preparer is not None and hasattr(preparer, "get_batch_kwargs"):
+            prepared = preparer.get_batch_kwargs(sequences, device=X.device)
+            input_ids = prepared["input_ids"]  # (b, l)
+            boundary_ids = preparer.boundary_token_ids.to(  # (n_boundary,)
+                device=input_ids.device, dtype=input_ids.dtype
+            )
+            # E1 wraps each raw sequence in BOS, context-label, terminal-label,
+            # and EOS tokens. Only amino-acid rows are biological residues.
+            M = M.to(dtype=torch.bool) & ~torch.isin(input_ids, boundary_ids)  # (b, l)
+        if need_attentions:
+            raise ValueError("parti is not available for tokenizer-free E1 embedding.")
+        return EmbeddingBatch(  # X: (b, l, d); residue_mask: (b, l)
+            X=X,
+            residue_mask=M.to(dtype=torch.bool),
+        )
+    if tokenizer is None:
+        raise ValueError("A tokenizer is required for this model's embedding path.")
+
+    tokenize_kwargs: dict[str, Any] = {
+        "return_tensors": "pt",
+        "padding": True,
+        "truncation": truncate,
+    }
+    if max_length is not None and truncate:
+        # ``max_length`` is a biological-residue limit. Tokenizer limits include
+        # boundary tokens, so reserve their declared width instead of dropping
+        # residues at the exact boundary.
+        special_token_count = 0
+        num_special_tokens_to_add = getattr(tokenizer, "num_special_tokens_to_add", None)
+        if callable(num_special_tokens_to_add):
+            special_token_count = int(num_special_tokens_to_add(pair=False))
+        tokenize_kwargs["max_length"] = max_length + special_token_count
+    sequence_tokenizer = getattr(model, "_tokenize_sequence_batch", None)
+    if callable(sequence_tokenizer):
+        encoded = sequence_tokenizer(sequences, tokenizer=tokenizer, **tokenize_kwargs)
+    else:
+        encoded = tokenizer(sequences, **tokenize_kwargs)
+    device = _model_device(model)
+    input_ids = encoded["input_ids"].to(device)  # (b, l)
+    attention_mask = encoded.get(  # (b, l)
+        "attention_mask",
+        input_ids.new_ones(input_ids.shape),
+    ).to(device)
+    M = _biological_residue_mask(input_ids, attention_mask, tokenizer)  # (b, l)
+    if need_attentions:
+        # Validate l before either the backbone or its quadratic attention graph
+        # is materialized. M has shape (b, l).
+        _validate_parti_length(M)
+    X = model._embed(input_ids, attention_mask, **model_kwargs)  # (b, l, d)
+    attentions = None
+    if need_attentions:
+        output = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_attentions=True,
+            return_dict=True,
+        )
+        attentions = getattr(output, "attentions", None)  # each: (b, h, l, l)
+        if attentions is None:
+            raise ValueError("The model did not return attentions required by parti.")
+    return EmbeddingBatch(  # X: (b, l, d); M: (b, l)
+        X=X,
+        residue_mask=M,
+        attentions=attentions,
+    )
+
+
+@dataclass(eq=False)
+class BatchExecutor:
+    """Model and batch policy for one bounded embedding window at a time."""
+
+    model: Any
+    batch_size: int
+    max_tokens_per_batch: int | None
+    max_length: int | None
+    truncate: bool
+    model_kwargs: dict[str, Any]
+    hidden_state_source: str
+    normalized_decoder_inputs: tuple[str, ...] | None
+    decoder_input_ids: Tensor | None
+    decoder_attention_mask: Tensor | None
+    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None
+    tokenizer: Any | None
+    store_all_hidden_states: bool
+    full_embeddings: bool
+    dtype: torch.dtype | None
+    pooler: Pooler | None
+    attention_backend: str | None
+    need_attentions: bool
+    model_type: str = field(init=False)
+    resolved_tokenizer: Any = field(init=False)
+
+    def __post_init__(self) -> None:
+        config = getattr(self.model, "config", None)
+        self.model_type = str(getattr(config, "model_type", "")).lower()
+        self.resolved_tokenizer = (
+            self.tokenizer if self.tokenizer is not None else getattr(self.model, "tokenizer", None)
+        )
+
+    def run_window(
+        self,
+        window_records: Sequence[EmbeddingInput],
+        *,
+        window_start: int,
+    ) -> tuple[list[EmbeddingRecord], dict[str, tuple[int, int]]]:
+        """Restore source order after length-bucketed inference and pooling."""
+
+        pool_slices: dict[str, tuple[int, int]] = {}
+        window_results: dict[int, EmbeddingRecord] = {}
+        for local_positions in _planned_batches(
+            window_records,
+            range(len(window_records)),
+            batch_size=self.batch_size,
+            max_tokens_per_batch=self.max_tokens_per_batch,
+            max_length=self.max_length,
+            truncate=self.truncate,
+        ):
+            batch_positions = [window_start + position for position in local_positions]
+            batch_records = [window_records[position] for position in local_positions]
+            sequences = [
+                record.sequence[: self.max_length]
+                if self.truncate and self.max_length is not None
+                else record.sequence
+                for record in batch_records
+            ]
+            batch_model_kwargs = dict(self.model_kwargs)
+            if self.model_type == "fast_ankh" or self.hidden_state_source == "decoder":
+                batch_model_kwargs["hidden_state_source"] = self.hidden_state_source
+                if self.normalized_decoder_inputs is not None:
+                    batch_model_kwargs["decoder_inputs"] = [
+                        self.normalized_decoder_inputs[position] for position in batch_positions
+                    ]
+                if self.decoder_input_ids is not None:
+                    # decoder_input_ids: (n_records, l_decoder)
+                    indices = torch.tensor(  # (b,)
+                        batch_positions,
+                        device=self.decoder_input_ids.device,
+                        dtype=torch.long,
+                    )
+                    batch_model_kwargs["decoder_input_ids"] = (  # (b, l_decoder)
+                        self.decoder_input_ids.index_select(0, indices)
+                    )
+                if self.decoder_attention_mask is not None:
+                    # decoder_attention_mask: (n_records, l_decoder)
+                    indices = torch.tensor(  # (b,)
+                        batch_positions,
+                        device=self.decoder_attention_mask.device,
+                        dtype=torch.long,
+                    )
+                    batch_model_kwargs["decoder_attention_mask"] = (
+                        self.decoder_attention_mask.index_select(0, indices)  # (b, l_decoder)
+                    )
+            custom_batch = self._embedding_batch_fn or getattr(self.model, "_embedding_batch", None)
+            if custom_batch is not None:
+                if self.model_type == "fast_ankh":
+                    batch = custom_batch(
+                        sequences,
+                        tokenizer=self.resolved_tokenizer,
+                        max_length=self.max_length,
+                        truncate=self.truncate,
+                        need_attentions=self.need_attentions,
+                        **batch_model_kwargs,
+                    )
+                else:
+                    batch = custom_batch(sequences, **batch_model_kwargs)
+                if not isinstance(batch, EmbeddingBatch):
+                    raise TypeError("_embedding_batch must return EmbeddingBatch.")
+            else:
+                batch = _generic_embedding_batch(
+                    self.model,
+                    sequences,
+                    tokenizer=self.tokenizer,
+                    max_length=self.max_length,
+                    truncate=self.truncate,
+                    need_attentions=self.need_attentions,
+                    model_kwargs=batch_model_kwargs,
+                )
+            X = batch.X  # (b, l, d) or (b, n_states, l, d)
+            raw_mask = batch.residue_mask  # (b, l)
+            if not isinstance(X, Tensor) or not isinstance(raw_mask, Tensor):
+                raise TypeError("Embedding batches must provide Tensor X and residue_mask.")
+            if X.is_meta or raw_mask.is_meta:
+                raise ValueError("Embedding batches cannot contain meta tensors.")
+            if not X.is_floating_point():
+                raise TypeError("Embedding batches must use a floating-point X dtype.")
+            if raw_mask.is_complex() or not bool(torch.isfinite(raw_mask).all()):
+                raise ValueError("Embedding residue_mask must contain finite binary values.")
+            if not bool(((raw_mask == 0) | (raw_mask == 1)).all()):
+                raise ValueError("Embedding residue_mask must contain finite binary values.")
+            M = raw_mask.to(device=X.device, dtype=torch.bool)  # (b, l)
+            valid_X_shape = (
+                X.ndim == 3
+                and X.shape[0] == len(batch_records)
+                and X.shape[-1] > 0
+                and M.shape == X.shape[:2]
+            )
+            valid_all_states_shape = (
+                X.ndim == 4
+                and self.store_all_hidden_states
+                and self.full_embeddings
+                and X.shape[0] == len(batch_records)
+                and X.shape[1] > 0
+                and X.shape[-1] > 0
+                and M.shape == (X.shape[0], X.shape[2])
+            )
+            if not (valid_X_shape or valid_all_states_shape):
+                raise ValueError(
+                    "Embedding batches must provide X with shape (b, l, d), or "
+                    "(b, states, l, d) when storing all hidden states, and "
+                    "residue_mask with shape (b, l)."
+                )
+            if not bool(M.any(dim=1).all()):
+                raise ValueError("Every embedding sample must contain a biological residue.")
+            finite_selected = (  # X.shape
+                torch.isfinite(X) | ~M.unsqueeze(-1)
+                if X.ndim == 3
+                else torch.isfinite(X) | ~M[:, None, :, None]
+            )
+            if not bool(finite_selected.all()):
+                raise ValueError("Biological residue embeddings produced non-finite output.")
+            if self.need_attentions:
+                # Validate the biological graph only after mask integrity is established.
+                _validate_parti_length(M)
+            if self.dtype is not None:
+                X = X.to(dtype=self.dtype)  # unchanged shape
+
+            if self.full_embeddings:
+                if X.ndim == 4:
+                    values = [
+                        X_i[:, M_i, :].detach().cpu()  # (n_states, r_i, d)
+                        for X_i, M_i in zip(X, M, strict=True)
+                    ]
+                else:
+                    values = _residue_embeddings(X, M)  # each: (r_i, d)
+            else:
+                if self.pooler is None:
+                    raise RuntimeError(
+                        "Pooled embedding output was requested without an initialized pooler."
+                    )
+                Y = self.pooler(  # (b, n_poolers * d)
+                    X,
+                    M,
+                    attentions=batch.attentions,
+                    attention_backend=self.attention_backend,
+                )
+                pool_slices = self.pooler.output_slices(X.shape[-1])
+                values = list(Y.detach().cpu().unbind(0))  # each: (n_poolers * d,)
+            for position, record, value in zip(batch_positions, batch_records, values, strict=True):
+                window_results[position] = EmbeddingRecord(record.id, record.sequence, value)
+
+        new_records = [
+            window_results[position]
+            for position in range(window_start, window_start + len(window_records))
+        ]
+        return new_records, pool_slices

fastplms/embeddings/identity.py

@@ -0,0 +1,510 @@
+"""Deterministic identity for embedding inputs, models, tokenizers, and execution."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import platform
+import torch
+from collections.abc import Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any
+from torch import Tensor
+
+from .inputs import _InputSpool
+from .storage import tensor_sha256
+from .types import EmbeddingInput
+
+
+_RUN_FINGERPRINT_SCHEMA_VERSION = 3
+_MODEL_STATE_HASH_CHUNK_BYTES = 16 * 1024**2
+
+
+def _model_device(model: Any) -> torch.device:
+    try:
+        return torch.device(next(model.parameters()).device)
+    except (AttributeError, StopIteration):
+        return torch.device("cpu")
+
+
+def _attention_backend(model: Any) -> str | None:
+    config = getattr(model, "config", None)
+    for name in ("_attn_implementation", "attn_implementation", "attn_backend"):
+        value = getattr(config, name, None)
+        if value:
+            return str(value)
+    return None
+
+
+def _attention_kernel_metadata(backend: str | None) -> dict[str, Any] | None:
+    if backend not in {"flash_attention_2", "flash_attention_3"}:
+        return None
+    from fastplms.registry import get_model_registry
+
+    spec = get_model_registry().attention_kernels[backend]
+    return {
+        "repository": spec.repository,
+        "revision": spec.revision,
+        "version": spec.version,
+        "expected_variant": spec.expected_variant,
+        "dtypes": list(spec.dtypes),
+    }
+
+
+def _fingerprint_jsonable(value: Any) -> Any:
+    if isinstance(value, Mapping):
+        return {str(key): _fingerprint_jsonable(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_fingerprint_jsonable(item) for item in value]
+    if isinstance(value, (set, frozenset)):
+        return sorted((_fingerprint_jsonable(item) for item in value), key=repr)
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, Tensor):
+        return {
+            "dtype": str(value.dtype).removeprefix("torch."),
+            "shape": list(value.shape),
+            "sha256": tensor_sha256(value),
+        }
+    if isinstance(value, torch.dtype):
+        return str(value).removeprefix("torch.")
+    if isinstance(value, torch.device):
+        return str(value)
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    return {
+        "class": f"{value.__class__.__module__}.{value.__class__.__qualname__}",
+        "value": str(value),
+    }
+
+
+def _tokenizer_content_sha256(tokenizer: Any) -> str:
+    content: dict[str, Any] = {
+        "init_kwargs": getattr(tokenizer, "init_kwargs", None),
+        "special_tokens_map": getattr(tokenizer, "special_tokens_map", None),
+        "model_max_length": getattr(tokenizer, "model_max_length", None),
+        "padding_side": getattr(tokenizer, "padding_side", None),
+        "truncation_side": getattr(tokenizer, "truncation_side", None),
+    }
+    get_vocab = getattr(tokenizer, "get_vocab", None)
+    if callable(get_vocab):
+        content["vocabulary"] = get_vocab()
+    get_added_vocab = getattr(tokenizer, "get_added_vocab", None)
+    if callable(get_added_vocab):
+        content["added_vocabulary"] = get_added_vocab()
+    backend = getattr(tokenizer, "backend_tokenizer", None)
+    backend_to_str = getattr(backend, "to_str", None)
+    if callable(backend_to_str):
+        content["backend"] = backend_to_str()
+    serialized = json.dumps(
+        _fingerprint_jsonable(content),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+    ).encode()
+    return hashlib.sha256(serialized).hexdigest()
+
+
+def _tokenizer_metadata(model: Any, tokenizer: Any | None) -> dict[str, Any]:
+    resolved = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None)
+    if resolved is None:
+        # Raw-sequence families such as E1 retain their loader context on the
+        # model/encoder rather than exposing a Transformers tokenizer. Bind the
+        # non-secret source policy to resume identity without serializing a Hub
+        # token or forcing lazy tokenizer initialization.
+        for candidate in (model, getattr(model, "model", None)):
+            settings = getattr(candidate, "__dict__", {}).get("_fastplms_tokenizer_kwargs")
+            if isinstance(settings, Mapping):
+                token_value = settings.get("token")
+                return {
+                    "mode": "native-sequence",
+                    "source": (
+                        str(settings.get("tokenizer_source"))
+                        if settings.get("tokenizer_source") is not None
+                        else None
+                    ),
+                    "revision": settings.get("revision"),
+                    "cache_dir": (
+                        str(settings.get("cache_dir"))
+                        if settings.get("cache_dir") is not None
+                        else None
+                    ),
+                    "local_files_only": bool(settings.get("local_files_only", False)),
+                    "token_policy": (
+                        "disabled"
+                        if token_value is False
+                        else "provided"
+                        if token_value is not None
+                        else "default"
+                    ),
+                }
+        return {"mode": "native-sequence"}
+    return {
+        "mode": "tokenizer",
+        "class": f"{resolved.__class__.__module__}.{resolved.__class__.__qualname__}",
+        "name_or_path": getattr(resolved, "name_or_path", None),
+        "vocab_size": getattr(resolved, "vocab_size", None),
+        "special_token_ids": list(getattr(resolved, "all_special_ids", ())),
+        "content_sha256": _tokenizer_content_sha256(resolved),
+    }
+
+
+def _software_versions() -> dict[str, str | None]:
+    try:
+        import fastplms
+
+        fastplms_version = fastplms.__version__
+    except (AttributeError, ImportError):
+        fastplms_version = None
+    try:
+        import safetensors
+
+        safetensors_version = safetensors.__version__
+    except ImportError:
+        safetensors_version = None
+    try:
+        import transformers
+
+        transformers_version = transformers.__version__
+    except ImportError:
+        transformers_version = None
+    return {
+        "fastplms": fastplms_version,
+        "python": platform.python_version(),
+        "safetensors": safetensors_version,
+        "torch": torch.__version__,
+        "torch_cuda": torch.version.cuda,
+        "transformers": transformers_version,
+    }
+
+
+def _adapter_identity_metadata(model: Any) -> dict[str, Any] | None:
+    """Return deterministic PEFT/adapter identity without tensor payloads."""
+
+    peft_config = getattr(model, "peft_config", None)
+    if not isinstance(peft_config, Mapping) or not peft_config:
+        return None
+    configurations: dict[str, Any] = {}
+    for name, config in sorted(peft_config.items(), key=lambda item: str(item[0])):
+        to_dict = getattr(config, "to_dict", None)
+        if callable(to_dict):
+            value = to_dict()
+        else:
+            try:
+                value = vars(config)
+            except TypeError:
+                value = config
+        configurations[str(name)] = _fingerprint_jsonable(value)
+    active_adapters = getattr(model, "active_adapters", None)
+    if callable(active_adapters):
+        active_adapters = active_adapters()
+    return {
+        "active": _fingerprint_jsonable(active_adapters),
+        "configurations": configurations,
+    }
+
+
+def _execution_identity_metadata(model: Any) -> dict[str, Any]:
+    """Capture runtime policy that can change persisted numerical results."""
+
+    parameter_dtypes = sorted(
+        {
+            str(parameter.dtype).removeprefix("torch.")
+            for parameter in getattr(model, "parameters", lambda: ())()
+        }
+    )
+    return {
+        "device": _model_device(model).type,
+        "hf_device_map": _fingerprint_jsonable(getattr(model, "hf_device_map", None)),
+        "parameter_dtypes": parameter_dtypes,
+        "software": _software_versions(),
+    }
+
+
+def _first_metadata_value(*values: Any) -> Any:
+    for value in values:
+        if isinstance(value, str):
+            if value.strip():
+                return value
+        elif value is not None:
+            return value
+    return None
+
+
+def _model_identity_metadata(model: Any) -> dict[str, Any]:
+    """Resolve model and checkpoint identity, including local artifact fallbacks."""
+
+    config = getattr(model, "config", None)
+    checkpoint_revision = _first_metadata_value(
+        getattr(config, "fastplms_checkpoint_revision", None),
+        getattr(config, "_commit_hash", None),
+    )
+    return {
+        "model_id": _first_metadata_value(
+            getattr(config, "fastplms_model_id", None),
+            getattr(config, "_name_or_path", None),
+        ),
+        "model_revision": _first_metadata_value(
+            getattr(config, "_commit_hash", None),
+            checkpoint_revision,
+        ),
+        "checkpoint_repo_id": getattr(config, "fastplms_checkpoint_repo_id", None),
+        "checkpoint_revision": checkpoint_revision,
+        "checkpoint_hash": _first_metadata_value(
+            getattr(model, "checkpoint_hash", None),
+            getattr(config, "checkpoint_hash", None),
+            getattr(config, "fastplms_checkpoint_hash", None),
+        ),
+        "weights_revision": getattr(config, "fastplms_weights_revision", None),
+        "runtime_revision": getattr(config, "fastplms_runtime_revision", None),
+        "source_tree_sha256": getattr(config, "fastplms_source_tree_sha256", None),
+        "runtime_bundle_sha256": getattr(config, "fastplms_runtime_bundle_sha256", None),
+    }
+
+
+def _bounded_tensor_chunks(X: Tensor, max_elements: int) -> Iterable[Tensor]:
+    """Yield X in logical row-major order without materializing a full copy."""
+
+    # X: (...)
+    if X.numel() == 0:
+        return
+    if X.ndim == 0:
+        yield X
+        return
+    trailing_elements = 1
+    for size in X.shape[1:]:
+        trailing_elements *= int(size)
+    if trailing_elements <= max_elements:
+        rows_per_chunk = max(1, max_elements // trailing_elements)
+        for start in range(0, X.shape[0], rows_per_chunk):
+            yield X[start : start + rows_per_chunk]  # (chunk_rows, ...)
+        return
+    for row in X:
+        yield from _bounded_tensor_chunks(row, max_elements)
+
+
+def _model_state_sha256(model: Any) -> str:
+    """Hash named parameters and persistent buffers using bounded CPU copies."""
+
+    # Never cache this digest from tensor identity or ``Tensor._version``.
+    # ``Parameter.data`` and independent tensor aliases can mutate shared storage
+    # without changing either signal, while persisted resume identity must bind
+    # the authoritative bytes visible at the start of this run.
+    state = model.state_dict(keep_vars=True)
+    digest = hashlib.sha256()
+    for name, value in sorted(state.items()):
+        if not isinstance(value, Tensor):
+            raise TypeError(f"Model state entry {name!r} is not a tensor.")
+        if value.is_meta:
+            raise ValueError(
+                f"Cannot fingerprint meta-device model state entry {name!r}; pass "
+                "model_state_fingerprint with a caller-owned state identity."
+            )
+        header = json.dumps(
+            {
+                "name": name,
+                "dtype": str(value.dtype).removeprefix("torch."),
+                "shape": list(value.shape),
+            },
+            sort_keys=True,
+            separators=(",", ":"),
+        ).encode()
+        digest.update(len(header).to_bytes(8, "big"))
+        digest.update(header)
+        max_elements = max(1, _MODEL_STATE_HASH_CHUNK_BYTES // value.element_size())
+        for chunk in _bounded_tensor_chunks(value.detach(), max_elements):
+            cpu_chunk = chunk.to(device="cpu").contiguous()  # chunk.shape
+            digest.update(cpu_chunk.reshape(-1).view(torch.uint8).numpy().tobytes())
+    return digest.hexdigest()
+
+
+def _input_sha256(records: Iterable[EmbeddingInput]) -> str:
+    """Hash an ordered input stream without constructing a duplicate JSON payload."""
+
+    precomputed = getattr(records, "input_fingerprint", None)
+    if isinstance(precomputed, str):
+        return precomputed
+    digest = hashlib.sha256()
+    count = 0
+    for record in records:
+        count += 1
+        for value in (record.id, record.sequence):
+            encoded = value.encode("utf-8")
+            digest.update(len(encoded).to_bytes(8, "big"))
+            digest.update(encoded)
+    digest.update(count.to_bytes(8, "big"))
+    return digest.hexdigest()
+
+
+def _run_fingerprint(
+    model: Any,
+    records: Sequence[EmbeddingInput],
+    *,
+    pooling: Sequence[str],
+    full_embeddings: bool,
+    max_length: int | None,
+    truncate: bool,
+    dtype: torch.dtype | None,
+    model_kwargs: dict[str, Any],
+    tokenizer_metadata: dict[str, Any],
+    model_state_fingerprint: str | None,
+    persist_output: bool,
+    embedding_context: Mapping[str, Any],
+    batch_size: int,
+    batch_window_size: int,
+    max_tokens_per_batch: int | None,
+) -> tuple[str, str, str | None, str]:
+    input_fingerprint = _input_sha256(records)
+    attention_backend = _attention_backend(model)
+    model_identity = _model_identity_metadata(model)
+    if model_state_fingerprint is None and persist_output:
+        resolved_model_state_fingerprint = _model_state_sha256(model)
+        model_state_fingerprint_source = "computed"
+    elif model_state_fingerprint is not None:
+        resolved_model_state_fingerprint = model_state_fingerprint.strip()
+        if not resolved_model_state_fingerprint:
+            raise ValueError("model_state_fingerprint must not be empty.")
+        model_state_fingerprint_source = "caller"
+    else:
+        resolved_model_state_fingerprint = None
+        model_state_fingerprint_source = "not-computed"
+    payload = {
+        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+        "input_fingerprint": input_fingerprint,
+        "model_state_fingerprint": resolved_model_state_fingerprint,
+        "model_state_fingerprint_source": model_state_fingerprint_source,
+        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
+        **model_identity,
+        "attention_backend": attention_backend,
+        "attention_kernel": _attention_kernel_metadata(attention_backend),
+        "layer": repr(
+            getattr(model, "embedding_layer", model_kwargs.get("hidden_state_index", -1))
+        ),
+        "projection": getattr(model, "embedding_projection", None),
+        "esmc_source": getattr(model, "_esmc_source", None),
+        "esmc_revision": getattr(model, "_esmc_source_revision", None),
+        "esmc_files": getattr(model, "_esmc_source_files", None),
+        "token_policy": getattr(model, "embedding_token_policy", None),
+        "tokenizer": tokenizer_metadata,
+        "adapter": _adapter_identity_metadata(model),
+        "execution": _execution_identity_metadata(model),
+        "embedding_context": _fingerprint_jsonable(embedding_context),
+        "pooling": list(pooling),
+        "full_embeddings": full_embeddings,
+        "max_length": max_length,
+        "truncate": truncate,
+        "dtype": str(dtype) if dtype is not None else None,
+        "batching": {
+            "batch_size": batch_size,
+            "batch_window_size": batch_window_size,
+            "max_tokens_per_batch": max_tokens_per_batch,
+            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
+        },
+        "model_kwargs": {
+            key: _fingerprint_jsonable(value) for key, value in sorted(model_kwargs.items())
+        },
+        "residue_mask_policy": "attention-mask-minus-special-tokens",
+    }
+    run_fingerprint = hashlib.sha256(
+        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
+    ).hexdigest()
+    return (
+        input_fingerprint,
+        run_fingerprint,
+        resolved_model_state_fingerprint,
+        model_state_fingerprint_source,
+    )
+
+
+def _ordered_string_sha256(values: Sequence[str]) -> str:
+    digest = hashlib.sha256()
+    for value in values:
+        encoded = value.encode("utf-8")
+        digest.update(len(encoded).to_bytes(8, "big"))
+        digest.update(encoded)
+    digest.update(len(values).to_bytes(8, "big"))
+    return digest.hexdigest()
+
+
+def _embedding_context(
+    model: Any,
+    records: Sequence[EmbeddingInput],
+    *,
+    hidden_state_source: str,
+    decoder_inputs: Sequence[str] | None,
+    decoder_input_ids: Tensor | None,
+    decoder_attention_mask: Tensor | None,
+    model_kwargs: Mapping[str, Any],
+) -> tuple[dict[str, Any], tuple[str, ...] | None]:
+    if hidden_state_source not in {"encoder", "decoder"}:
+        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
+    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
+    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
+        raise TypeError("hidden_state_index must be an integer.")
+    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
+    if not isinstance(store_all_hidden_states, bool):
+        raise TypeError("store_all_hidden_states must be a boolean.")
+    normalized_decoder_inputs: tuple[str, ...] | None = None
+    has_decoder_inputs = decoder_inputs is not None
+    has_decoder_ids = decoder_input_ids is not None
+    if hidden_state_source == "encoder":
+        if has_decoder_inputs or has_decoder_ids or decoder_attention_mask is not None:
+            raise ValueError("Decoder inputs are only valid when hidden_state_source='decoder'.")
+    else:
+        if has_decoder_inputs == has_decoder_ids:
+            raise ValueError(
+                "Decoder embedding requires exactly one of decoder_inputs or decoder_input_ids."
+            )
+    decoder_input_fingerprint: str | None = None
+    if decoder_inputs is not None:
+        if isinstance(decoder_inputs, (str, bytes)) or not isinstance(decoder_inputs, Sequence):
+            raise TypeError("decoder_inputs must be an aligned sequence of strings.")
+        normalized_decoder_inputs = tuple(decoder_inputs)
+        if not all(isinstance(value, str) and value for value in normalized_decoder_inputs):
+            raise ValueError("decoder_inputs must contain non-empty strings.")
+        if len(normalized_decoder_inputs) != len(records):
+            raise ValueError("decoder_inputs must align one-to-one with embedding inputs.")
+        decoder_input_fingerprint = _ordered_string_sha256(normalized_decoder_inputs)
+        if decoder_attention_mask is not None:
+            raise ValueError("decoder_attention_mask requires decoder_input_ids.")
+    if decoder_input_ids is not None:
+        if not isinstance(decoder_input_ids, Tensor) or decoder_input_ids.ndim != 2:
+            raise ValueError("decoder_input_ids must have shape (batch, sequence).")
+        if decoder_input_ids.shape[0] != len(records):
+            raise ValueError("decoder_input_ids must align one-to-one with embedding inputs.")
+        if decoder_input_ids.dtype == torch.bool or decoder_input_ids.is_floating_point():
+            raise TypeError("decoder_input_ids must use an integer token dtype.")
+        decoder_input_fingerprint = tensor_sha256(decoder_input_ids)
+    decoder_mask_fingerprint: str | None = None
+    if decoder_attention_mask is not None:
+        if not isinstance(decoder_attention_mask, Tensor):
+            raise TypeError("decoder_attention_mask must be a tensor.")
+        if decoder_input_ids is None or decoder_attention_mask.shape != decoder_input_ids.shape:
+            raise ValueError("decoder_attention_mask must match decoder_input_ids shape.")
+        decoder_mask_fingerprint = tensor_sha256(decoder_attention_mask)
+
+    context: dict[str, Any] = {
+        "hidden_state_source": hidden_state_source,
+        "hidden_state_index": hidden_state_index,
+        "store_all_hidden_states": store_all_hidden_states,
+        "decoder_input_fingerprint": decoder_input_fingerprint,
+        "decoder_attention_mask_fingerprint": decoder_mask_fingerprint,
+        "decoder_alignment": "input-position" if hidden_state_source == "decoder" else None,
+    }
+    metadata_hook = getattr(model, "_embedding_metadata", None)
+    model_metadata: Mapping[str, Any] | None = None
+    if callable(metadata_hook):
+        model_metadata = metadata_hook(**context)
+        if not isinstance(model_metadata, Mapping):
+            raise TypeError("_embedding_metadata must return a mapping.")
+        context["model_embedding"] = _fingerprint_jsonable(model_metadata)
+    if hidden_state_source == "decoder":
+        has_decoder_batch = callable(getattr(model, "_embedding_batch", None))
+        declares_decoder_stack = (
+            model_metadata is not None and model_metadata.get("hidden_state_stack") == "decoder"
+        )
+        if not has_decoder_batch or not declares_decoder_stack:
+            raise ValueError(
+                f"{model.__class__.__name__} does not declare decoder embedding support."
+            )
+    return context, normalized_decoder_inputs

fastplms/embeddings/inputs.py

@@ -0,0 +1,263 @@
+"""Normalize ordered inputs and plan bounded windows without retaining a full stream."""
+
+from __future__ import annotations
+
+import hashlib
+import sqlite3
+import tempfile
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from pathlib import Path
+from typing import overload
+
+from .types import EmbeddingInput
+
+
+def iter_fasta(path: str | Path) -> Iterator[EmbeddingInput]:
+    """Yield FASTA records in source order without reading the file into memory."""
+
+    identifier: str | None = None
+    sequence_parts: list[str] = []
+    found_record = False
+    with Path(path).open("r", encoding="utf-8") as handle:
+        for line_number, raw_line in enumerate(handle, start=1):
+            line = raw_line.strip()
+            if not line:
+                continue
+            if line.startswith(">"):
+                if identifier is not None:
+                    found_record = True
+                    yield EmbeddingInput(identifier, "".join(sequence_parts))
+                identifier = line[1:].strip().split(maxsplit=1)[0]
+                if not identifier:
+                    raise ValueError(f"Missing FASTA identifier on line {line_number}.")
+                sequence_parts = []
+            else:
+                if identifier is None:
+                    raise ValueError(
+                        f"Sequence data precedes the first FASTA header on line {line_number}."
+                    )
+                sequence_parts.append("".join(line.split()))
+    if identifier is not None:
+        found_record = True
+        yield EmbeddingInput(identifier, "".join(sequence_parts))
+    if not found_record:
+        raise ValueError(f"No FASTA records found in {path}.")
+
+
+def parse_fasta(path: str | Path) -> list[EmbeddingInput]:
+    """Parse FASTA records while preserving identifiers, order, and duplicates."""
+
+    return list(iter_fasta(path))
+
+
+def _normalize_input_item(
+    position: int,
+    item: str | EmbeddingInput | tuple[str, str],
+) -> EmbeddingInput:
+    if isinstance(item, EmbeddingInput):
+        return item
+    if isinstance(item, str):
+        return EmbeddingInput(str(position), item)
+    if isinstance(item, tuple) and len(item) == 2:
+        return EmbeddingInput(str(item[0]), str(item[1]))
+    raise TypeError(
+        "inputs must contain sequences, EmbeddingInput values, or (id, sequence) tuples."
+    )
+
+
+class _InputSpool(Sequence[EmbeddingInput]):
+    """Immutable disk-backed normalized inputs with an incremental digest."""
+
+    def __init__(
+        self,
+        values: Iterable[str | EmbeddingInput | tuple[str, str]],
+    ) -> None:
+        self._temporary: tempfile.TemporaryDirectory[str] | None = tempfile.TemporaryDirectory(
+            prefix="fastplms-inputs-"
+        )
+        self.path = Path(self._temporary.name) / "inputs.sqlite"
+        self._connection: sqlite3.Connection | None = sqlite3.connect(self.path)
+        self._connection.execute(
+            "CREATE TABLE inputs ("
+            "position INTEGER PRIMARY KEY, input_id TEXT NOT NULL, sequence TEXT NOT NULL)"
+        )
+        digest = hashlib.sha256()
+        count = 0
+        pending: list[tuple[int, str, str]] = []
+        try:
+            for position, item in enumerate(values):
+                record = _normalize_input_item(position, item)
+                for value in (record.id, record.sequence):
+                    encoded = value.encode("utf-8")
+                    digest.update(len(encoded).to_bytes(8, "big"))
+                    digest.update(encoded)
+                pending.append((position, record.id, record.sequence))
+                count += 1
+                if len(pending) == 1_024:
+                    self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
+                    pending.clear()
+            if pending:
+                self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
+            if count == 0:
+                raise ValueError("inputs must contain at least one sequence.")
+            self._connection.commit()
+            self._connection.close()
+            self._connection = sqlite3.connect(
+                f"{self.path.resolve().as_uri()}?mode=ro",
+                uri=True,
+            )
+        except BaseException:
+            self.close()
+            raise
+        digest.update(count.to_bytes(8, "big"))
+        self.input_fingerprint = digest.hexdigest()
+        self._count = count
+
+    def _require_connection(self) -> sqlite3.Connection:
+        if self._connection is None:
+            raise RuntimeError("Input spool is closed.")
+        return self._connection
+
+    def __len__(self) -> int:
+        return self._count
+
+    def __iter__(self) -> Iterator[EmbeddingInput]:
+        cursor = self._require_connection().execute(
+            "SELECT input_id, sequence FROM inputs ORDER BY position"
+        )
+        while rows := cursor.fetchmany(1_024):
+            for input_id, sequence in rows:
+                yield EmbeddingInput(input_id, sequence)
+
+    @overload
+    def __getitem__(self, index: int, /) -> EmbeddingInput: ...
+
+    @overload
+    def __getitem__(self, index: slice, /) -> list[EmbeddingInput]: ...
+
+    def __getitem__(self, index: int | slice) -> EmbeddingInput | list[EmbeddingInput]:
+        connection = self._require_connection()
+
+        if isinstance(index, slice):
+            start, stop, step = index.indices(self._count)
+            if step != 1:
+                return [self[position] for position in range(start, stop, step)]
+            rows = connection.execute(
+                "SELECT input_id, sequence FROM inputs "
+                "WHERE position >= ? AND position < ? ORDER BY position",
+                (start, stop),
+            ).fetchall()
+            return [EmbeddingInput(input_id, sequence) for input_id, sequence in rows]
+        position = index + self._count if index < 0 else index
+        if position < 0 or position >= self._count:
+            raise IndexError(index)
+        row = connection.execute(
+            "SELECT input_id, sequence FROM inputs WHERE position = ?", (position,)
+        ).fetchone()
+        if row is None:
+            raise IndexError(index)
+        return EmbeddingInput(row[0], row[1])
+
+    def close(self) -> None:
+        connection = getattr(self, "_connection", None)
+        if connection is not None:
+            connection.close()
+            self._connection = None
+        temporary = getattr(self, "_temporary", None)
+        if temporary is not None:
+            temporary.cleanup()
+            self._temporary = None
+
+    def __del__(self) -> None:
+        self.close()
+
+
+def _normalize_inputs(
+    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
+    *,
+    disk_backed: bool,
+) -> Sequence[EmbeddingInput]:
+    is_fasta_path = isinstance(inputs, Path)
+    if isinstance(inputs, str):
+        try:
+            is_fasta_path = Path(inputs).is_file()
+        except OSError:
+            is_fasta_path = False
+    should_spool = disk_backed or is_fasta_path or not isinstance(inputs, (str, Sequence, Mapping))
+    values: Iterable[str | EmbeddingInput | tuple[str, str]]
+    if isinstance(inputs, Path):
+        values = iter_fasta(inputs)
+    elif isinstance(inputs, str):
+        values = iter_fasta(inputs) if is_fasta_path else [inputs]
+    elif isinstance(inputs, Mapping):
+        values = inputs.items()
+    else:
+        values = inputs
+    if should_spool:
+        return _InputSpool(values)
+    records: list[EmbeddingInput] = []
+    for position, item in enumerate(values):
+        records.append(_normalize_input_item(position, item))
+    if not records:
+        raise ValueError("inputs must contain at least one sequence.")
+    return records
+
+
+def _validate_untruncated_lengths(
+    records: Sequence[EmbeddingInput],
+    *,
+    max_length: int | None,
+    truncate: bool,
+) -> None:
+    """Fail before inference when a biological-residue limit would be exceeded."""
+
+    if max_length is None or truncate:
+        return
+    for position, record in enumerate(records):
+        residue_count = len(record.sequence)
+        if residue_count > max_length:
+            raise ValueError(
+                f"Input at position {position} with id {record.id!r} has "
+                f"{residue_count} biological residues, exceeding max_length={max_length} "
+                "while truncate=False."
+            )
+
+
+def _planned_batches(
+    records: Sequence[EmbeddingInput],
+    positions: range,
+    *,
+    batch_size: int,
+    max_tokens_per_batch: int | None,
+    max_length: int | None,
+    truncate: bool,
+) -> Iterator[list[int]]:
+    """Length-bucket one bounded window while retaining stable output positions."""
+
+    def effective_length(position: int) -> int:
+        length = len(records[position].sequence)
+        return min(length, max_length) if truncate and max_length is not None else length
+
+    ordered = sorted(positions, key=lambda position: (-effective_length(position), position))
+    batch: list[int] = []
+    longest = 0
+    for position in ordered:
+        length = effective_length(position)
+        if max_tokens_per_batch is not None and length > max_tokens_per_batch:
+            raise ValueError(
+                f"Input at position {position} has {length} residues, exceeding "
+                f"max_tokens_per_batch={max_tokens_per_batch}."
+            )
+        candidate_longest = max(longest, length)
+        exceeds_tokens = (
+            max_tokens_per_batch is not None
+            and candidate_longest * (len(batch) + 1) > max_tokens_per_batch
+        )
+        if batch and (len(batch) >= batch_size or exceeds_tokens):
+            yield batch
+            batch = []
+            longest = 0
+        batch.append(position)
+        longest = max(longest, length)
+    if batch:
+        yield batch

fastplms/embeddings/output.py

@@ -0,0 +1,215 @@
+"""Resume validation and transactional publication of ordered embedding windows."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from .identity import _RUN_FINGERPRINT_SCHEMA_VERSION
+from .pooling import Pooler
+from .storage import (
+    SafetensorsStreamWriter,
+    append_sqlite_records,
+    initialize_sqlite_run,
+    load_result,
+    load_sqlite_result,
+    safetensors_result_exists,
+    save_result,
+    tensor_sha256,
+    update_sqlite_run_metadata,
+)
+from .types import EmbeddingInput, EmbeddingRecord, EmbeddingResult, LazyTensorReference
+
+
+def _output_exists(path: str | Path, format: str) -> bool:
+    path = Path(path)
+    if format == "sqlite":
+        return path.is_file()
+    return safetensors_result_exists(path)
+
+
+def _output_descriptor(position: int, record: EmbeddingRecord) -> dict[str, Any]:
+    tensor = record.tensor
+    if isinstance(tensor, LazyTensorReference):
+        dtype = tensor.dtype
+        shape = tensor.shape
+        digest = tensor.sha256
+    else:
+        dtype = str(tensor.dtype).removeprefix("torch.")
+        shape = tuple(tensor.shape)
+        digest = tensor_sha256(tensor)
+    return {
+        "position": position,
+        "id": record.id,
+        "dtype": dtype,
+        "shape": shape,
+        "sha256": digest,
+    }
+
+
+class EmbeddingOutput:
+    """Own the resumable prefix and the commit state of one output destination."""
+
+    def __init__(
+        self,
+        records: Sequence[EmbeddingInput],
+        *,
+        output: str | Path | None,
+        format: str,
+        resume: bool,
+        shard_size: int,
+        run_fingerprint: str,
+        input_fingerprint: str,
+        model_state_fingerprint: str | None,
+        model_state_fingerprint_source: str,
+        pooler: Pooler | None,
+        pooling_names: Sequence[str],
+    ) -> None:
+        self.output = output
+        self.format = format
+        self.shard_size = shard_size
+        self.completed: EmbeddingResult | None = None
+        output_already_exists = output is not None and _output_exists(output, format)
+        existing: EmbeddingResult | None = None
+        self.start_position = 0
+        if output is not None and resume and output_already_exists:
+            if format == "sqlite":
+                try:
+                    existing = load_sqlite_result(output, run_id=run_fingerprint)
+                except KeyError:
+                    existing = load_result(output, format=format)
+            else:
+                existing = load_result(output, format=format)
+            if existing.metadata.get("fingerprint_schema_version") != (
+                _RUN_FINGERPRINT_SCHEMA_VERSION
+            ):
+                raise ValueError(
+                    "Existing embeddings use an incompatible run fingerprint schema; "
+                    "choose another output or set resume=False."
+                )
+            if existing.metadata.get("run_fingerprint") != run_fingerprint:
+                raise ValueError(
+                    "Existing embeddings were produced by a different run fingerprint; "
+                    "choose another output or set resume=False."
+                )
+            if len(existing) > len(records):
+                raise ValueError(
+                    "Existing embeddings are not an ordered prefix of the requested inputs."
+                )
+            prefix_matches = all(
+                (observed.id, observed.sequence) == (expected.id, expected.sequence)
+                for expected, observed in zip(records, existing, strict=False)
+            )
+            if not prefix_matches:
+                raise ValueError(
+                    "Existing embeddings are not an ordered prefix of the requested inputs."
+                )
+            if len(existing) == len(records) and existing.metadata.get("complete", True):
+                self.completed = existing
+                return
+            self.start_position = len(existing)
+
+        self.sqlite_run_id: str | None = None
+        self.sqlite_replace_on_first_commit = False
+        self.sqlite_initial_metadata: dict[str, Any] | None = None
+        if output is not None and format == "sqlite":
+            self.sqlite_initial_metadata = {
+                "format_version": 1,
+                "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+                "run_fingerprint": run_fingerprint,
+                "input_fingerprint": input_fingerprint,
+                "model_state_fingerprint": model_state_fingerprint,
+                "model_state_fingerprint_source": model_state_fingerprint_source,
+                "complete": False,
+            }
+            self.sqlite_run_id = run_fingerprint
+            if not resume and output_already_exists:
+                try:
+                    load_sqlite_result(output, run_id=run_fingerprint)
+                except KeyError:
+                    pass
+                else:
+                    # Keep an exact prior run readable until replacement inference
+                    # has produced the first complete commit window.
+                    self.sqlite_replace_on_first_commit = True
+            if not self.sqlite_replace_on_first_commit:
+                initialize_sqlite_run(
+                    output,
+                    self.sqlite_initial_metadata,
+                    resume=resume,
+                )
+
+        stream_safetensors = output is not None and format == "safetensors"
+        self.output_records: list[EmbeddingRecord] = (
+            [] if self.sqlite_run_id is not None or stream_safetensors else list(existing or ())
+        )
+        self.output_descriptors: list[dict[str, Any]] | None = [] if output is None else None
+        self.pool_slices: dict[str, tuple[int, int]] = {}
+        if existing and pooler is not None:
+            pooled_width = existing[0].load_tensor().shape[-1]
+            if pooled_width % len(pooling_names) != 0:
+                raise ValueError("Stored pooled width is inconsistent with pooling metadata.")
+            self.pool_slices = pooler.output_slices(pooled_width // len(pooling_names))
+
+        self.safetensors_writer: SafetensorsStreamWriter | None = None
+        if stream_safetensors:
+            if output is None:
+                raise RuntimeError(
+                    "Safetensors streaming was enabled without an output destination."
+                )
+            transactional_overwrite = output_already_exists and not resume
+            self.safetensors_writer = SafetensorsStreamWriter(
+                output,
+                {
+                    "format_version": 1,
+                    "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+                    "run_fingerprint": run_fingerprint,
+                    "input_fingerprint": input_fingerprint,
+                    "model_state_fingerprint": model_state_fingerprint,
+                    "model_state_fingerprint_source": model_state_fingerprint_source,
+                    "complete": False,
+                },
+                shard_size=shard_size,
+                existing=existing or (),
+                reuse_existing=bool(resume and existing is not None),
+                publish_initial=not transactional_overwrite,
+                publish_incremental=not transactional_overwrite,
+            )
+
+    def append(self, window_start: int, new_records: list[EmbeddingRecord]) -> None:
+        """Commit a complete ordered window at the storage format's granularity."""
+
+        if self.output_descriptors is not None:
+            self.output_descriptors.extend(
+                _output_descriptor(window_start + offset, record)
+                for offset, record in enumerate(new_records)
+            )
+        if self.output is not None and self.sqlite_run_id is not None:
+            append_sqlite_records(
+                self.output,
+                self.sqlite_run_id,
+                window_start,
+                new_records,
+                replace_metadata=(
+                    self.sqlite_initial_metadata if self.sqlite_replace_on_first_commit else None
+                ),
+            )
+            self.sqlite_replace_on_first_commit = False
+        elif self.safetensors_writer is not None:
+            self.safetensors_writer.append(new_records)
+        else:
+            self.output_records.extend(new_records)
+
+    def finish(self, metadata: dict[str, Any]) -> EmbeddingResult:
+        """Publish completion only after every window has committed."""
+
+        if self.output is not None and self.sqlite_run_id is not None:
+            update_sqlite_run_metadata(self.output, self.sqlite_run_id, metadata)
+            return load_sqlite_result(self.output, run_id=self.sqlite_run_id)
+        if self.safetensors_writer is not None:
+            return self.safetensors_writer.publish(complete=True, metadata=metadata)
+        result = EmbeddingResult(self.output_records, metadata)
+        if self.output is not None:
+            return save_result(result, self.output, format=self.format, shard_size=self.shard_size)
+        return result

fastplms/embeddings/runner.py

@@ -1,1583 +1,425 @@
-"""Model-independent dataset embedding orchestration."""
-
-from __future__ import annotations
-
-import hashlib
-import json
-import platform
-import sqlite3
-import tempfile
-import torch
-from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
-from contextlib import contextmanager
-from pathlib import Path
-from typing import Any, overload
+"""Coordinate input preparation, run identity, batch execution, and publication."""
+
+from __future__ import annotations
+
+import torch
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any
 from torch import Tensor
 
-from .pooling import Pooler
-from .storage import (
-    SafetensorsStreamWriter,
-    append_sqlite_records,
-    initialize_sqlite_run,
-    load_result,
-    load_sqlite_result,
-    safetensors_result_exists,
-    save_result,
-    tensor_sha256,
-    update_sqlite_run_metadata,
-)
-from .types import (
-    EmbeddingBatch,
-    EmbeddingInput,
-    EmbeddingRecord,
-    EmbeddingResult,
-    LazyTensorReference,
-)
-
-
-_MAX_PARTI_RESIDUES = 2_048
-_RUN_FINGERPRINT_SCHEMA_VERSION = 3
-_MODEL_STATE_HASH_CHUNK_BYTES = 16 * 1024**2
-_DEFAULT_BATCH_WINDOW_MULTIPLIER = 16
-_SUPPORTED_STORAGE_FORMATS = frozenset({"safetensors", "sqlite"})
-
-
-def _validate_parti_length(M: Tensor) -> None:
-    """Reject an oversized attention graph before model inference."""
-
-    # M: (b, l)
-    n_residues = int(M.to(dtype=torch.int64).sum(dim=1).max().item())
-    if n_residues > _MAX_PARTI_RESIDUES:
-        raise ValueError(f"parti supports at most {_MAX_PARTI_RESIDUES:,} biological residues.")
-
-
-def select_hidden_state_embeddings(
-    last_hidden_state: Tensor,
-    hidden_states: tuple[Tensor, ...] | None,
-    *,
-    hidden_state_index: int = -1,
-    store_all_hidden_states: bool = False,
-) -> Tensor:
-    """Select one hidden state or stack every state without changing values."""
-    # last_hidden_state and each hidden_states entry: (b, l, d)
-    if store_all_hidden_states:
-        if not hidden_states:
-            raise ValueError("store_all_hidden_states requires model hidden states.")
-        # H has shape (b, n, l, d), where n follows the model's output order.
-        return torch.stack(hidden_states, dim=1)  # (b, n, l, d)
-    if hidden_state_index == -1:
-        return last_hidden_state  # (b, l, d)
-    if not hidden_states:
-        raise ValueError("hidden_state_index requires model hidden states.")
-    return hidden_states[hidden_state_index]  # (b, l, d)
-
-
-def iter_fasta(path: str | Path) -> Iterator[EmbeddingInput]:
-    """Yield FASTA records in source order without reading the file into memory."""
-
-    identifier: str | None = None
-    sequence_parts: list[str] = []
-    found_record = False
-    with Path(path).open("r", encoding="utf-8") as handle:
-        for line_number, raw_line in enumerate(handle, start=1):
-            line = raw_line.strip()
-            if not line:
-                continue
-            if line.startswith(">"):
-                if identifier is not None:
-                    found_record = True
-                    yield EmbeddingInput(identifier, "".join(sequence_parts))
-                identifier = line[1:].strip().split(maxsplit=1)[0]
-                if not identifier:
-                    raise ValueError(f"Missing FASTA identifier on line {line_number}.")
-                sequence_parts = []
-            else:
-                if identifier is None:
-                    raise ValueError(
-                        f"Sequence data precedes the first FASTA header on line {line_number}."
-                    )
-                sequence_parts.append("".join(line.split()))
-    if identifier is not None:
-        found_record = True
-        yield EmbeddingInput(identifier, "".join(sequence_parts))
-    if not found_record:
-        raise ValueError(f"No FASTA records found in {path}.")
-
-
-def parse_fasta(path: str | Path) -> list[EmbeddingInput]:
-    """Parse FASTA records while preserving identifiers, order, and duplicates."""
-
-    return list(iter_fasta(path))
-
-
-def _normalize_input_item(
-    position: int,
-    item: str | EmbeddingInput | tuple[str, str],
-) -> EmbeddingInput:
-    if isinstance(item, EmbeddingInput):
-        return item
-    if isinstance(item, str):
-        return EmbeddingInput(str(position), item)
-    if isinstance(item, tuple) and len(item) == 2:
-        return EmbeddingInput(str(item[0]), str(item[1]))
-    raise TypeError(
-        "inputs must contain sequences, EmbeddingInput values, or (id, sequence) tuples."
-    )
-
-
-class _InputSpool(Sequence[EmbeddingInput]):
-    """Immutable disk-backed normalized inputs with an incremental digest."""
-
-    def __init__(
-        self,
-        values: Iterable[str | EmbeddingInput | tuple[str, str]],
-    ) -> None:
-        self._temporary: tempfile.TemporaryDirectory[str] | None = tempfile.TemporaryDirectory(
-            prefix="fastplms-inputs-"
-        )
-        self.path = Path(self._temporary.name) / "inputs.sqlite"
-        self._connection: sqlite3.Connection | None = sqlite3.connect(self.path)
-        self._connection.execute(
-            "CREATE TABLE inputs ("
-            "position INTEGER PRIMARY KEY, input_id TEXT NOT NULL, sequence TEXT NOT NULL)"
-        )
-        digest = hashlib.sha256()
-        count = 0
-        pending: list[tuple[int, str, str]] = []
-        try:
-            for position, item in enumerate(values):
-                record = _normalize_input_item(position, item)
-                for value in (record.id, record.sequence):
-                    encoded = value.encode("utf-8")
-                    digest.update(len(encoded).to_bytes(8, "big"))
-                    digest.update(encoded)
-                pending.append((position, record.id, record.sequence))
-                count += 1
-                if len(pending) == 1_024:
-                    self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
-                    pending.clear()
-            if pending:
-                self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
-            if count == 0:
-                raise ValueError("inputs must contain at least one sequence.")
-            self._connection.commit()
-            self._connection.close()
-            self._connection = sqlite3.connect(
-                f"{self.path.resolve().as_uri()}?mode=ro",
-                uri=True,
-            )
-        except BaseException:
-            self.close()
-            raise
-        digest.update(count.to_bytes(8, "big"))
-        self.input_fingerprint = digest.hexdigest()
-        self._count = count
-
-    def _require_connection(self) -> sqlite3.Connection:
-        if self._connection is None:
-            raise RuntimeError("Input spool is closed.")
-        return self._connection
-
-    def __len__(self) -> int:
-        return self._count
-
-    def __iter__(self) -> Iterator[EmbeddingInput]:
-        cursor = self._require_connection().execute(
-            "SELECT input_id, sequence FROM inputs ORDER BY position"
-        )
-        while rows := cursor.fetchmany(1_024):
-            for input_id, sequence in rows:
-                yield EmbeddingInput(input_id, sequence)
-
-    @overload
-    def __getitem__(self, index: int, /) -> EmbeddingInput: ...
-
-    @overload
-    def __getitem__(self, index: slice, /) -> list[EmbeddingInput]: ...
-
-    def __getitem__(self, index: int | slice) -> EmbeddingInput | list[EmbeddingInput]:
-        connection = self._require_connection()
-
-        if isinstance(index, slice):
-            start, stop, step = index.indices(self._count)
-            if step != 1:
-                return [self[position] for position in range(start, stop, step)]
-            rows = connection.execute(
-                "SELECT input_id, sequence FROM inputs "
-                "WHERE position >= ? AND position < ? ORDER BY position",
-                (start, stop),
-            ).fetchall()
-            return [EmbeddingInput(input_id, sequence) for input_id, sequence in rows]
-        position = index + self._count if index < 0 else index
-        if position < 0 or position >= self._count:
-            raise IndexError(index)
-        row = connection.execute(
-            "SELECT input_id, sequence FROM inputs WHERE position = ?", (position,)
-        ).fetchone()
-        if row is None:
-            raise IndexError(index)
-        return EmbeddingInput(row[0], row[1])
-
-    def close(self) -> None:
-        connection = getattr(self, "_connection", None)
-        if connection is not None:
-            connection.close()
-            self._connection = None
-        temporary = getattr(self, "_temporary", None)
-        if temporary is not None:
-            temporary.cleanup()
-            self._temporary = None
-
-    def __del__(self) -> None:
-        self.close()
-
-
-def _normalize_inputs(
-    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
-    *,
-    disk_backed: bool,
-) -> Sequence[EmbeddingInput]:
-    is_fasta_path = isinstance(inputs, Path)
-    if isinstance(inputs, str):
-        try:
-            is_fasta_path = Path(inputs).is_file()
-        except OSError:
-            is_fasta_path = False
-    should_spool = disk_backed or is_fasta_path or not isinstance(inputs, (str, Sequence, Mapping))
-    values: Iterable[str | EmbeddingInput | tuple[str, str]]
-    if isinstance(inputs, Path):
-        values = iter_fasta(inputs)
-    elif isinstance(inputs, str):
-        values = iter_fasta(inputs) if is_fasta_path else [inputs]
-    elif isinstance(inputs, Mapping):
-        values = inputs.items()
-    else:
-        values = inputs
-    if should_spool:
-        return _InputSpool(values)
-    records: list[EmbeddingInput] = []
-    for position, item in enumerate(values):
-        records.append(_normalize_input_item(position, item))
-    if not records:
-        raise ValueError("inputs must contain at least one sequence.")
-    return records
-
-
-def _validate_untruncated_lengths(
-    records: Sequence[EmbeddingInput],
-    *,
-    max_length: int | None,
-    truncate: bool,
-) -> None:
-    """Fail before inference when a biological-residue limit would be exceeded."""
-
-    if max_length is None or truncate:
-        return
-    for position, record in enumerate(records):
-        residue_count = len(record.sequence)
-        if residue_count > max_length:
-            raise ValueError(
-                f"Input at position {position} with id {record.id!r} has "
-                f"{residue_count} biological residues, exceeding max_length={max_length} "
-                "while truncate=False."
-            )
-
-
-def _model_device(model: Any) -> torch.device:
-    try:
-        return torch.device(next(model.parameters()).device)
-    except (AttributeError, StopIteration):
-        return torch.device("cpu")
-
-
-def _attention_backend(model: Any) -> str | None:
-    config = getattr(model, "config", None)
-    for name in ("_attn_implementation", "attn_implementation", "attn_backend"):
-        value = getattr(config, name, None)
-        if value:
-            return str(value)
-    return None
-
-
-def _attention_kernel_metadata(backend: str | None) -> dict[str, Any] | None:
-    if backend not in {"flash_attention_2", "flash_attention_3"}:
-        return None
-    from fastplms.registry import get_model_registry
-
-    spec = get_model_registry().attention_kernels[backend]
-    return {
-        "repository": spec.repository,
-        "revision": spec.revision,
-        "version": spec.version,
-        "expected_variant": spec.expected_variant,
-        "dtypes": list(spec.dtypes),
-    }
-
-
-def _fingerprint_jsonable(value: Any) -> Any:
-    if isinstance(value, Mapping):
-        return {str(key): _fingerprint_jsonable(item) for key, item in value.items()}
-    if isinstance(value, (list, tuple)):
-        return [_fingerprint_jsonable(item) for item in value]
-    if isinstance(value, (set, frozenset)):
-        return sorted((_fingerprint_jsonable(item) for item in value), key=repr)
-    if isinstance(value, Path):
-        return str(value)
-    if isinstance(value, Tensor):
-        return {
-            "dtype": str(value.dtype).removeprefix("torch."),
-            "shape": list(value.shape),
-            "sha256": tensor_sha256(value),
-        }
-    if isinstance(value, torch.dtype):
-        return str(value).removeprefix("torch.")
-    if isinstance(value, torch.device):
-        return str(value)
-    if value is None or isinstance(value, (str, int, float, bool)):
-        return value
-    return {
-        "class": f"{value.__class__.__module__}.{value.__class__.__qualname__}",
-        "value": str(value),
-    }
-
-
-def _tokenizer_content_sha256(tokenizer: Any) -> str:
-    content: dict[str, Any] = {
-        "init_kwargs": getattr(tokenizer, "init_kwargs", None),
-        "special_tokens_map": getattr(tokenizer, "special_tokens_map", None),
-        "model_max_length": getattr(tokenizer, "model_max_length", None),
-        "padding_side": getattr(tokenizer, "padding_side", None),
-        "truncation_side": getattr(tokenizer, "truncation_side", None),
-    }
-    get_vocab = getattr(tokenizer, "get_vocab", None)
-    if callable(get_vocab):
-        content["vocabulary"] = get_vocab()
-    get_added_vocab = getattr(tokenizer, "get_added_vocab", None)
-    if callable(get_added_vocab):
-        content["added_vocabulary"] = get_added_vocab()
-    backend = getattr(tokenizer, "backend_tokenizer", None)
-    backend_to_str = getattr(backend, "to_str", None)
-    if callable(backend_to_str):
-        content["backend"] = backend_to_str()
-    serialized = json.dumps(
-        _fingerprint_jsonable(content),
-        sort_keys=True,
-        separators=(",", ":"),
-        ensure_ascii=False,
-    ).encode()
-    return hashlib.sha256(serialized).hexdigest()
-
-
-def _tokenizer_metadata(model: Any, tokenizer: Any | None) -> dict[str, Any]:
-    resolved = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None)
-    if resolved is None:
-        # Raw-sequence families such as E1 retain their loader context on the
-        # model/encoder rather than exposing a Transformers tokenizer. Bind the
-        # non-secret source policy to resume identity without serializing a Hub
-        # token or forcing lazy tokenizer initialization.
-        for candidate in (model, getattr(model, "model", None)):
-            settings = getattr(candidate, "__dict__", {}).get("_fastplms_tokenizer_kwargs")
-            if isinstance(settings, Mapping):
-                token_value = settings.get("token")
-                return {
-                    "mode": "native-sequence",
-                    "source": (
-                        str(settings.get("tokenizer_source"))
-                        if settings.get("tokenizer_source") is not None
-                        else None
-                    ),
-                    "revision": settings.get("revision"),
-                    "cache_dir": (
-                        str(settings.get("cache_dir"))
-                        if settings.get("cache_dir") is not None
-                        else None
-                    ),
-                    "local_files_only": bool(settings.get("local_files_only", False)),
-                    "token_policy": (
-                        "disabled"
-                        if token_value is False
-                        else "provided"
-                        if token_value is not None
-                        else "default"
-                    ),
-                }
-        return {"mode": "native-sequence"}
-    return {
-        "mode": "tokenizer",
-        "class": f"{resolved.__class__.__module__}.{resolved.__class__.__qualname__}",
-        "name_or_path": getattr(resolved, "name_or_path", None),
-        "vocab_size": getattr(resolved, "vocab_size", None),
-        "special_token_ids": list(getattr(resolved, "all_special_ids", ())),
-        "content_sha256": _tokenizer_content_sha256(resolved),
-    }
-
-
-@contextmanager
-def _temporary_eval(model: Any) -> Iterator[None]:
-    was_training = getattr(model, "training", None)
-    eval_method = getattr(model, "eval", None)
-    train_method = getattr(model, "train", None)
-    if (
-        not isinstance(was_training, bool)
-        or not callable(eval_method)
-        or not callable(train_method)
-    ):
-        yield
-        return
-    eval_method()
-    try:
-        yield
-    finally:
-        train_method(was_training)
-
-
-def _software_versions() -> dict[str, str | None]:
-    try:
-        import fastplms
-
-        fastplms_version = fastplms.__version__
-    except (AttributeError, ImportError):
-        fastplms_version = None
-    try:
-        import safetensors
-
-        safetensors_version = safetensors.__version__
-    except ImportError:
-        safetensors_version = None
-    try:
-        import transformers
-
-        transformers_version = transformers.__version__
-    except ImportError:
-        transformers_version = None
-    return {
-        "fastplms": fastplms_version,
-        "python": platform.python_version(),
-        "safetensors": safetensors_version,
-        "torch": torch.__version__,
-        "torch_cuda": torch.version.cuda,
-        "transformers": transformers_version,
-    }
-
-
-def _adapter_identity_metadata(model: Any) -> dict[str, Any] | None:
-    """Return deterministic PEFT/adapter identity without tensor payloads."""
-
-    peft_config = getattr(model, "peft_config", None)
-    if not isinstance(peft_config, Mapping) or not peft_config:
-        return None
-    configurations: dict[str, Any] = {}
-    for name, config in sorted(peft_config.items(), key=lambda item: str(item[0])):
-        to_dict = getattr(config, "to_dict", None)
-        if callable(to_dict):
-            value = to_dict()
-        else:
-            try:
-                value = vars(config)
-            except TypeError:
-                value = config
-        configurations[str(name)] = _fingerprint_jsonable(value)
-    active_adapters = getattr(model, "active_adapters", None)
-    if callable(active_adapters):
-        active_adapters = active_adapters()
-    return {
-        "active": _fingerprint_jsonable(active_adapters),
-        "configurations": configurations,
-    }
-
-
-def _execution_identity_metadata(model: Any) -> dict[str, Any]:
-    """Capture runtime policy that can change persisted numerical results."""
-
-    parameter_dtypes = sorted(
-        {
-            str(parameter.dtype).removeprefix("torch.")
-            for parameter in getattr(model, "parameters", lambda: ())()
-        }
-    )
-    return {
-        "device": _model_device(model).type,
-        "hf_device_map": _fingerprint_jsonable(getattr(model, "hf_device_map", None)),
-        "parameter_dtypes": parameter_dtypes,
-        "software": _software_versions(),
-    }
-
-
-def _biological_residue_mask(
-    input_ids: Tensor,
-    attention_mask: Tensor,
-    tokenizer: Any,
-) -> Tensor:
-    """Remove padding and tokenizer-declared special tokens from M."""
-
-    # input_ids, attention_mask: (b, l)
-    M = attention_mask.to(dtype=torch.bool)  # (b, l)
-    special_ids = tuple(int(token_id) for token_id in getattr(tokenizer, "all_special_ids", ()))
-    if special_ids:
-        specials = torch.tensor(  # (n_special,)
-            special_ids,
-            device=input_ids.device,
-            dtype=input_ids.dtype,
-        )
-        M = M & ~torch.isin(input_ids, specials)  # (b, l)
-    return M  # (b, l)
-
-
-def _generic_embedding_batch(
-    model: Any,
-    sequences: list[str],
-    *,
-    tokenizer: Any | None,
-    max_length: int | None,
-    truncate: bool,
-    need_attentions: bool,
-    model_kwargs: dict[str, Any],
-) -> EmbeddingBatch:
-    config = getattr(model, "config", None)
-    model_type = str(getattr(config, "model_type", "")).lower()
-    if tokenizer is None:
-        tokenizer = getattr(model, "tokenizer", None)
-
-    if tokenizer is None and model_type == "e1":
-        output = model._embed(sequences, return_attention_mask=True, **model_kwargs)
-        if not isinstance(output, tuple) or len(output) != 2:
-            raise TypeError("E1 _embed must return (X, residue_mask).")
-        X, M = output  # (b, l, d), (b, l)
-        preparer = getattr(model, "prep_tokens", None)
-        if preparer is not None and hasattr(preparer, "get_batch_kwargs"):
-            prepared = preparer.get_batch_kwargs(sequences, device=X.device)
-            input_ids = prepared["input_ids"]  # (b, l)
-            boundary_ids = preparer.boundary_token_ids.to(  # (n_boundary,)
-                device=input_ids.device, dtype=input_ids.dtype
-            )
-            # E1 wraps each raw sequence in BOS, context-label, terminal-label,
-            # and EOS tokens. Only amino-acid rows are biological residues.
-            M = M.to(dtype=torch.bool) & ~torch.isin(input_ids, boundary_ids)  # (b, l)
-        if need_attentions:
-            raise ValueError("parti is not available for tokenizer-free E1 embedding.")
-        return EmbeddingBatch(  # X: (b, l, d); residue_mask: (b, l)
-            X=X,
-            residue_mask=M.to(dtype=torch.bool),
-        )
-    if tokenizer is None:
-        raise ValueError("A tokenizer is required for this model's embedding path.")
-
-    tokenize_kwargs: dict[str, Any] = {
-        "return_tensors": "pt",
-        "padding": True,
-        "truncation": truncate,
-    }
-    if max_length is not None and truncate:
-        # ``max_length`` is a biological-residue limit. Tokenizer limits include
-        # boundary tokens, so reserve their declared width instead of dropping
-        # residues at the exact boundary.
-        special_token_count = 0
-        num_special_tokens_to_add = getattr(tokenizer, "num_special_tokens_to_add", None)
-        if callable(num_special_tokens_to_add):
-            special_token_count = int(num_special_tokens_to_add(pair=False))
-        tokenize_kwargs["max_length"] = max_length + special_token_count
-    sequence_tokenizer = getattr(model, "_tokenize_sequence_batch", None)
-    if callable(sequence_tokenizer):
-        encoded = sequence_tokenizer(sequences, tokenizer=tokenizer, **tokenize_kwargs)
-    else:
-        encoded = tokenizer(sequences, **tokenize_kwargs)
-    device = _model_device(model)
-    input_ids = encoded["input_ids"].to(device)  # (b, l)
-    attention_mask = encoded.get(  # (b, l)
-        "attention_mask",
-        input_ids.new_ones(input_ids.shape),
-    ).to(device)
-    M = _biological_residue_mask(input_ids, attention_mask, tokenizer)  # (b, l)
-    if need_attentions:
-        # Validate l before either the backbone or its quadratic attention graph
-        # is materialized. M has shape (b, l).
-        _validate_parti_length(M)
-    X = model._embed(input_ids, attention_mask, **model_kwargs)  # (b, l, d)
-    attentions = None
-    if need_attentions:
-        output = model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            output_attentions=True,
-            return_dict=True,
-        )
-        attentions = getattr(output, "attentions", None)  # each: (b, h, l, l)
-        if attentions is None:
-            raise ValueError("The model did not return attentions required by parti.")
-    return EmbeddingBatch(  # X: (b, l, d); M: (b, l)
-        X=X,
-        residue_mask=M,
-        attentions=attentions,
-    )
-
-
-def _first_metadata_value(*values: Any) -> Any:
-    for value in values:
-        if isinstance(value, str):
-            if value.strip():
-                return value
-        elif value is not None:
-            return value
-    return None
-
-
-def _model_identity_metadata(model: Any) -> dict[str, Any]:
-    """Resolve model and checkpoint identity, including local artifact fallbacks."""
-
-    config = getattr(model, "config", None)
-    checkpoint_revision = _first_metadata_value(
-        getattr(config, "fastplms_checkpoint_revision", None),
-        getattr(config, "_commit_hash", None),
-    )
-    return {
-        "model_id": _first_metadata_value(
-            getattr(config, "fastplms_model_id", None),
-            getattr(config, "_name_or_path", None),
-        ),
-        "model_revision": _first_metadata_value(
-            getattr(config, "_commit_hash", None),
-            checkpoint_revision,
-        ),
-        "checkpoint_repo_id": getattr(config, "fastplms_checkpoint_repo_id", None),
-        "checkpoint_revision": checkpoint_revision,
-        "checkpoint_hash": _first_metadata_value(
-            getattr(model, "checkpoint_hash", None),
-            getattr(config, "checkpoint_hash", None),
-            getattr(config, "fastplms_checkpoint_hash", None),
-        ),
-        "weights_revision": getattr(config, "fastplms_weights_revision", None),
-        "runtime_revision": getattr(config, "fastplms_runtime_revision", None),
-        "source_tree_sha256": getattr(config, "fastplms_source_tree_sha256", None),
-        "runtime_bundle_sha256": getattr(config, "fastplms_runtime_bundle_sha256", None),
-    }
-
-
-def _bounded_tensor_chunks(X: Tensor, max_elements: int) -> Iterable[Tensor]:
-    """Yield X in logical row-major order without materializing a full copy."""
-
-    # X: (...)
-    if X.numel() == 0:
-        return
-    if X.ndim == 0:
-        yield X
-        return
-    trailing_elements = 1
-    for size in X.shape[1:]:
-        trailing_elements *= int(size)
-    if trailing_elements <= max_elements:
-        rows_per_chunk = max(1, max_elements // trailing_elements)
-        for start in range(0, X.shape[0], rows_per_chunk):
-            yield X[start : start + rows_per_chunk]  # (chunk_rows, ...)
-        return
-    for row in X:
-        yield from _bounded_tensor_chunks(row, max_elements)
-
-
-def _model_state_sha256(model: Any) -> str:
-    """Hash named parameters and persistent buffers using bounded CPU copies."""
-
-    # Never cache this digest from tensor identity or ``Tensor._version``.
-    # ``Parameter.data`` and independent tensor aliases can mutate shared storage
-    # without changing either signal, while persisted resume identity must bind
-    # the authoritative bytes visible at the start of this run.
-    state = model.state_dict(keep_vars=True)
-    digest = hashlib.sha256()
-    for name, value in sorted(state.items()):
-        if not isinstance(value, Tensor):
-            raise TypeError(f"Model state entry {name!r} is not a tensor.")
-        if value.is_meta:
-            raise ValueError(
-                f"Cannot fingerprint meta-device model state entry {name!r}; pass "
-                "model_state_fingerprint with a caller-owned state identity."
-            )
-        header = json.dumps(
-            {
-                "name": name,
-                "dtype": str(value.dtype).removeprefix("torch."),
-                "shape": list(value.shape),
-            },
-            sort_keys=True,
-            separators=(",", ":"),
-        ).encode()
-        digest.update(len(header).to_bytes(8, "big"))
-        digest.update(header)
-        max_elements = max(1, _MODEL_STATE_HASH_CHUNK_BYTES // value.element_size())
-        for chunk in _bounded_tensor_chunks(value.detach(), max_elements):
-            cpu_chunk = chunk.to(device="cpu").contiguous()  # chunk.shape
-            digest.update(cpu_chunk.reshape(-1).view(torch.uint8).numpy().tobytes())
-    return digest.hexdigest()
-
-
-def _input_sha256(records: Iterable[EmbeddingInput]) -> str:
-    """Hash an ordered input stream without constructing a duplicate JSON payload."""
-
-    precomputed = getattr(records, "input_fingerprint", None)
-    if isinstance(precomputed, str):
-        return precomputed
-    digest = hashlib.sha256()
-    count = 0
-    for record in records:
-        count += 1
-        for value in (record.id, record.sequence):
-            encoded = value.encode("utf-8")
-            digest.update(len(encoded).to_bytes(8, "big"))
-            digest.update(encoded)
-    digest.update(count.to_bytes(8, "big"))
-    return digest.hexdigest()
-
-
-def _run_fingerprint(
-    model: Any,
-    records: Sequence[EmbeddingInput],
-    *,
-    pooling: Sequence[str],
-    full_embeddings: bool,
-    max_length: int | None,
-    truncate: bool,
-    dtype: torch.dtype | None,
-    model_kwargs: dict[str, Any],
-    tokenizer_metadata: dict[str, Any],
-    model_state_fingerprint: str | None,
-    persist_output: bool,
-    embedding_context: Mapping[str, Any],
-    batch_size: int,
-    batch_window_size: int,
-    max_tokens_per_batch: int | None,
-) -> tuple[str, str, str | None, str]:
-    input_fingerprint = _input_sha256(records)
-    attention_backend = _attention_backend(model)
-    model_identity = _model_identity_metadata(model)
-    if model_state_fingerprint is None and persist_output:
-        resolved_model_state_fingerprint = _model_state_sha256(model)
-        model_state_fingerprint_source = "computed"
-    elif model_state_fingerprint is not None:
-        resolved_model_state_fingerprint = model_state_fingerprint.strip()
-        if not resolved_model_state_fingerprint:
-            raise ValueError("model_state_fingerprint must not be empty.")
-        model_state_fingerprint_source = "caller"
-    else:
-        resolved_model_state_fingerprint = None
-        model_state_fingerprint_source = "not-computed"
-    payload = {
-        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-        "input_fingerprint": input_fingerprint,
-        "model_state_fingerprint": resolved_model_state_fingerprint,
-        "model_state_fingerprint_source": model_state_fingerprint_source,
-        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
-        **model_identity,
-        "attention_backend": attention_backend,
-        "attention_kernel": _attention_kernel_metadata(attention_backend),
-        "layer": repr(
-            getattr(model, "embedding_layer", model_kwargs.get("hidden_state_index", -1))
-        ),
-        "projection": getattr(model, "embedding_projection", None),
-        "esmc_source": getattr(model, "_esmc_source", None),
-        "esmc_revision": getattr(model, "_esmc_source_revision", None),
-        "esmc_files": getattr(model, "_esmc_source_files", None),
-        "token_policy": getattr(model, "embedding_token_policy", None),
-        "tokenizer": tokenizer_metadata,
-        "adapter": _adapter_identity_metadata(model),
-        "execution": _execution_identity_metadata(model),
-        "embedding_context": _fingerprint_jsonable(embedding_context),
-        "pooling": list(pooling),
-        "full_embeddings": full_embeddings,
-        "max_length": max_length,
-        "truncate": truncate,
-        "dtype": str(dtype) if dtype is not None else None,
-        "batching": {
-            "batch_size": batch_size,
-            "batch_window_size": batch_window_size,
-            "max_tokens_per_batch": max_tokens_per_batch,
-            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
-        },
-        "model_kwargs": {
-            key: _fingerprint_jsonable(value) for key, value in sorted(model_kwargs.items())
-        },
-        "residue_mask_policy": "attention-mask-minus-special-tokens",
-    }
-    run_fingerprint = hashlib.sha256(
-        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
-    ).hexdigest()
-    return (
-        input_fingerprint,
-        run_fingerprint,
-        resolved_model_state_fingerprint,
-        model_state_fingerprint_source,
-    )
-
-
-def _output_exists(path: str | Path, format: str) -> bool:
-    path = Path(path)
-    if format == "sqlite":
-        return path.is_file()
-    return safetensors_result_exists(path)
-
-
-def _output_descriptor(position: int, record: EmbeddingRecord) -> dict[str, Any]:
-    tensor = record.tensor
-    if isinstance(tensor, LazyTensorReference):
-        dtype = tensor.dtype
-        shape = tensor.shape
-        digest = tensor.sha256
-    else:
-        dtype = str(tensor.dtype).removeprefix("torch.")
-        shape = tuple(tensor.shape)
-        digest = tensor_sha256(tensor)
-    return {
-        "position": position,
-        "id": record.id,
-        "dtype": dtype,
-        "shape": shape,
-        "sha256": digest,
-    }
-
-
-def _ordered_string_sha256(values: Sequence[str]) -> str:
-    digest = hashlib.sha256()
-    for value in values:
-        encoded = value.encode("utf-8")
-        digest.update(len(encoded).to_bytes(8, "big"))
-        digest.update(encoded)
-    digest.update(len(values).to_bytes(8, "big"))
-    return digest.hexdigest()
-
-
-def _embedding_context(
-    model: Any,
-    records: Sequence[EmbeddingInput],
-    *,
-    hidden_state_source: str,
-    decoder_inputs: Sequence[str] | None,
-    decoder_input_ids: Tensor | None,
-    decoder_attention_mask: Tensor | None,
-    model_kwargs: Mapping[str, Any],
-) -> tuple[dict[str, Any], tuple[str, ...] | None]:
-    if hidden_state_source not in {"encoder", "decoder"}:
-        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
-    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
-    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
-        raise TypeError("hidden_state_index must be an integer.")
-    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
-    if not isinstance(store_all_hidden_states, bool):
-        raise TypeError("store_all_hidden_states must be a boolean.")
-    normalized_decoder_inputs: tuple[str, ...] | None = None
-    has_decoder_inputs = decoder_inputs is not None
-    has_decoder_ids = decoder_input_ids is not None
-    if hidden_state_source == "encoder":
-        if has_decoder_inputs or has_decoder_ids or decoder_attention_mask is not None:
-            raise ValueError("Decoder inputs are only valid when hidden_state_source='decoder'.")
-    else:
-        if has_decoder_inputs == has_decoder_ids:
-            raise ValueError(
-                "Decoder embedding requires exactly one of decoder_inputs or decoder_input_ids."
-            )
-    decoder_input_fingerprint: str | None = None
-    if decoder_inputs is not None:
-        if isinstance(decoder_inputs, (str, bytes)) or not isinstance(decoder_inputs, Sequence):
-            raise TypeError("decoder_inputs must be an aligned sequence of strings.")
-        normalized_decoder_inputs = tuple(decoder_inputs)
-        if not all(isinstance(value, str) and value for value in normalized_decoder_inputs):
-            raise ValueError("decoder_inputs must contain non-empty strings.")
-        if len(normalized_decoder_inputs) != len(records):
-            raise ValueError("decoder_inputs must align one-to-one with embedding inputs.")
-        decoder_input_fingerprint = _ordered_string_sha256(normalized_decoder_inputs)
-        if decoder_attention_mask is not None:
-            raise ValueError("decoder_attention_mask requires decoder_input_ids.")
-    if decoder_input_ids is not None:
-        if not isinstance(decoder_input_ids, Tensor) or decoder_input_ids.ndim != 2:
-            raise ValueError("decoder_input_ids must have shape (batch, sequence).")
-        if decoder_input_ids.shape[0] != len(records):
-            raise ValueError("decoder_input_ids must align one-to-one with embedding inputs.")
-        if decoder_input_ids.dtype == torch.bool or decoder_input_ids.is_floating_point():
-            raise TypeError("decoder_input_ids must use an integer token dtype.")
-        decoder_input_fingerprint = tensor_sha256(decoder_input_ids)
-    decoder_mask_fingerprint: str | None = None
-    if decoder_attention_mask is not None:
-        if not isinstance(decoder_attention_mask, Tensor):
-            raise TypeError("decoder_attention_mask must be a tensor.")
-        if decoder_input_ids is None or decoder_attention_mask.shape != decoder_input_ids.shape:
-            raise ValueError("decoder_attention_mask must match decoder_input_ids shape.")
-        decoder_mask_fingerprint = tensor_sha256(decoder_attention_mask)
-
-    context: dict[str, Any] = {
-        "hidden_state_source": hidden_state_source,
-        "hidden_state_index": hidden_state_index,
-        "store_all_hidden_states": store_all_hidden_states,
-        "decoder_input_fingerprint": decoder_input_fingerprint,
-        "decoder_attention_mask_fingerprint": decoder_mask_fingerprint,
-        "decoder_alignment": "input-position" if hidden_state_source == "decoder" else None,
-    }
-    metadata_hook = getattr(model, "_embedding_metadata", None)
-    model_metadata: Mapping[str, Any] | None = None
-    if callable(metadata_hook):
-        model_metadata = metadata_hook(**context)
-        if not isinstance(model_metadata, Mapping):
-            raise TypeError("_embedding_metadata must return a mapping.")
-        context["model_embedding"] = _fingerprint_jsonable(model_metadata)
-    if hidden_state_source == "decoder":
-        has_decoder_batch = callable(getattr(model, "_embedding_batch", None))
-        declares_decoder_stack = (
-            model_metadata is not None and model_metadata.get("hidden_state_stack") == "decoder"
-        )
-        if not has_decoder_batch or not declares_decoder_stack:
-            raise ValueError(
-                f"{model.__class__.__name__} does not declare decoder embedding support."
-            )
-    return context, normalized_decoder_inputs
-
-
-def _planned_batches(
-    records: Sequence[EmbeddingInput],
-    positions: range,
-    *,
-    batch_size: int,
-    max_tokens_per_batch: int | None,
-    max_length: int | None,
-    truncate: bool,
-) -> Iterator[list[int]]:
-    """Length-bucket one bounded window while retaining stable output positions."""
-
-    def effective_length(position: int) -> int:
-        length = len(records[position].sequence)
-        return min(length, max_length) if truncate and max_length is not None else length
-
-    ordered = sorted(positions, key=lambda position: (-effective_length(position), position))
-    batch: list[int] = []
-    longest = 0
-    for position in ordered:
-        length = effective_length(position)
-        if max_tokens_per_batch is not None and length > max_tokens_per_batch:
-            raise ValueError(
-                f"Input at position {position} has {length} residues, exceeding "
-                f"max_tokens_per_batch={max_tokens_per_batch}."
-            )
-        candidate_longest = max(longest, length)
-        exceeds_tokens = (
-            max_tokens_per_batch is not None
-            and candidate_longest * (len(batch) + 1) > max_tokens_per_batch
-        )
-        if batch and (len(batch) >= batch_size or exceeds_tokens):
-            yield batch
-            batch = []
-            longest = 0
-        batch.append(position)
-        longest = max(longest, length)
-    if batch:
-        yield batch
-
-
-def embed_dataset(
-    model: Any,
-    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
-    *,
-    batch_size: int = 2,
-    pooling: str | Sequence[str] | None = None,
-    full_embeddings: bool = False,
-    output: str | Path | None = None,
-    format: str = "safetensors",
-    resume: bool = True,
-    tokenizer: Any | None = None,
-    max_length: int | None = None,
-    truncate: bool = True,
-    dtype: torch.dtype | None = torch.float32,
-    shard_size: int = 2 * 1024**3,
-    model_state_fingerprint: str | None = None,
-    batch_window_size: int | None = None,
-    max_tokens_per_batch: int | None = None,
-    hidden_state_source: str = "encoder",
-    decoder_inputs: Sequence[str] | None = None,
-    decoder_input_ids: Tensor | None = None,
-    decoder_attention_mask: Tensor | None = None,
-    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None = None,
-    _embedding_batch_identity: Mapping[str, Any] | None = None,
-    _allowed_unsupported_pooling: Sequence[str] = (),
-    **model_kwargs: Any,
-) -> EmbeddingResult:
-    """Embed protein sequences with stable ordering and residue-only pooling."""
-
-    for name, value in (
-        ("batch_size", batch_size),
-        ("shard_size", shard_size),
-    ):
-        if not isinstance(value, int) or isinstance(value, bool):
-            raise TypeError(f"{name} must be a positive integer.")
-        if value <= 0:
-            raise ValueError(f"{name} must be a positive integer.")
-    for optional_name, optional_value in (
-        ("max_length", max_length),
-        ("max_tokens_per_batch", max_tokens_per_batch),
-        ("batch_window_size", batch_window_size),
-    ):
-        if optional_value is not None and (
-            not isinstance(optional_value, int) or isinstance(optional_value, bool)
-        ):
-            raise TypeError(f"{optional_name} must be a positive integer when provided.")
-        if optional_value is not None and optional_value <= 0:
-            raise ValueError(f"{optional_name} must be a positive integer when provided.")
-    for name, value in (
-        ("full_embeddings", full_embeddings),
-        ("resume", resume),
-        ("truncate", truncate),
-    ):
-        if not isinstance(value, bool):
-            raise TypeError(f"{name} must be a boolean.")
-    if not isinstance(format, str):
-        raise TypeError("format must be a string.")
-    if output is not None and not isinstance(output, (str, Path)):
-        raise TypeError("output must be a path or None.")
-    if model_state_fingerprint is not None and (
-        not isinstance(model_state_fingerprint, str) or not model_state_fingerprint
-    ):
-        raise ValueError("model_state_fingerprint must be a non-empty string when provided.")
-    if hidden_state_source not in {"encoder", "decoder"}:
-        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
-    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
-    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
-        raise TypeError("hidden_state_index must be an integer.")
-    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
-    if not isinstance(store_all_hidden_states, bool):
-        raise TypeError("store_all_hidden_states must be a boolean.")
-    if decoder_input_ids is not None:
-        if not isinstance(decoder_input_ids, Tensor):
-            raise TypeError("decoder_input_ids must be a tensor.")
-        if decoder_input_ids.is_meta:
-            raise ValueError("decoder_input_ids cannot be a meta tensor.")
-        if decoder_input_ids.ndim != 2 or decoder_input_ids.shape[1] == 0:
-            raise ValueError("decoder_input_ids must have non-empty shape (batch, sequence).")
-        if decoder_input_ids.dtype not in {torch.int32, torch.int64}:
-            raise TypeError("decoder_input_ids must use torch.int32 or torch.int64.")
-    if decoder_attention_mask is not None:
-        if not isinstance(decoder_attention_mask, Tensor):
-            raise TypeError("decoder_attention_mask must be a tensor.")
-        if decoder_attention_mask.is_meta:
-            raise ValueError("decoder_attention_mask cannot be a meta tensor.")
-        if decoder_attention_mask.is_complex() or not bool(
-            torch.isfinite(decoder_attention_mask).all()
-        ):
-            raise ValueError("decoder_attention_mask must contain finite binary values.")
-        if not bool(((decoder_attention_mask == 0) | (decoder_attention_mask == 1)).all()):
-            raise ValueError("decoder_attention_mask must contain finite binary values.")
-    pooling_names = (
-        (("mean",) if not full_embeddings else ())
-        if pooling is None
-        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
-    )
-    if full_embeddings and pooling is not None:
-        raise ValueError("full_embeddings=True cannot be combined with pooling.")
-    if not full_embeddings and not pooling_names:
-        raise ValueError("pooling is required unless full_embeddings=True.")
-    pooler = Pooler(pooling_names) if pooling_names else None
-
-    if batch_size <= 0:
-        raise ValueError("batch_size must be positive.")
-    if format == "pth" or (output is not None and Path(output).suffix.lower() == ".pth"):
-        raise ValueError("Writing pickle-based .pth embeddings is not supported.")
-    if format not in _SUPPORTED_STORAGE_FORMATS:
-        raise ValueError("format must be 'safetensors' or 'sqlite'.")
-    if max_length is not None and max_length <= 0:
-        raise ValueError("max_length must be positive when provided.")
-    if max_tokens_per_batch is not None and max_tokens_per_batch <= 0:
-        raise ValueError("max_tokens_per_batch must be positive when provided.")
-    if not isinstance(dtype, (torch.dtype, type(None))):
-        raise TypeError("dtype must be a torch.dtype or None.")
-    if batch_window_size is not None and batch_window_size <= 0:
-        raise ValueError("batch_window_size must be positive when provided.")
-    if _embedding_batch_fn is not None and not callable(_embedding_batch_fn):
-        raise TypeError("_embedding_batch_fn must be callable when provided.")
-    if _embedding_batch_fn is not None and _embedding_batch_identity is None:
-        raise ValueError(
-            "_embedding_batch_identity is required with _embedding_batch_fn so persisted "
-            "runs bind the family-specific embedding behavior."
-        )
-    if _embedding_batch_identity is not None and not isinstance(_embedding_batch_identity, Mapping):
-        raise TypeError("_embedding_batch_identity must be a mapping when provided.")
-    if isinstance(_allowed_unsupported_pooling, (str, bytes)) or not isinstance(
-        _allowed_unsupported_pooling, Sequence
-    ):
-        raise TypeError("_allowed_unsupported_pooling must be a sequence of pooler names.")
-    if not all(isinstance(name, str) for name in _allowed_unsupported_pooling):
-        raise TypeError("_allowed_unsupported_pooling must contain only strings.")
-    allowed_unsupported_pooling = frozenset(_allowed_unsupported_pooling)
-    if allowed_unsupported_pooling and _embedding_batch_fn is None:
-        raise ValueError(
-            "_allowed_unsupported_pooling is only valid with a family-specific _embedding_batch_fn."
-        )
-    resolved_batch_window_size = (
-        batch_size * _DEFAULT_BATCH_WINDOW_MULTIPLIER
-        if batch_window_size is None
-        else batch_window_size
-    )
-    if resolved_batch_window_size < batch_size:
-        raise ValueError("batch_window_size must be at least batch_size.")
-    records = _normalize_inputs(inputs, disk_backed=output is not None)
-    _validate_untruncated_lengths(
-        records,
-        max_length=max_length,
-        truncate=truncate,
-    )
-    pooling_names = (
-        (("mean",) if not full_embeddings else ())
-        if pooling is None
-        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
-    )
-    if full_embeddings:
-        if pooling is not None:
-            raise ValueError("full_embeddings=True cannot be combined with pooling.")
-    elif not pooling_names:
-        raise ValueError("pooling is required unless full_embeddings=True.")
-    store_all_hidden_states = bool(model_kwargs.get("store_all_hidden_states", False))
-    if store_all_hidden_states and not full_embeddings:
-        raise ValueError("store_all_hidden_states=True requires full_embeddings=True.")
-
-    unsupported = set(getattr(model, "embedding_unsupported_pooling", ()))
-    unknown_pooling_overrides = allowed_unsupported_pooling.difference(unsupported)
-    if unknown_pooling_overrides:
-        raise ValueError(
-            "_allowed_unsupported_pooling may only override poolers declared unsupported "
-            f"by the model; unknown overrides: {sorted(unknown_pooling_overrides)}."
-        )
-    unsupported.difference_update(allowed_unsupported_pooling)
-    requested_unsupported = unsupported.intersection(pooling_names)
-    if requested_unsupported:
-        raise ValueError(
-            f"{model.__class__.__name__} does not support pooling operations "
-            f"{sorted(requested_unsupported)}."
-        )
-
-    # Constructing the pooler validates names and duplicate operations before
-    # any checkpoint hashing, tokenization, or inference occurs.
-    pooler = Pooler(pooling_names) if pooling_names else None
-    embedding_context, normalized_decoder_inputs = _embedding_context(
-        model,
-        records,
-        hidden_state_source=hidden_state_source,
-        decoder_inputs=decoder_inputs,
-        decoder_input_ids=decoder_input_ids,
-        decoder_attention_mask=decoder_attention_mask,
-        model_kwargs=model_kwargs,
-    )
-    if _embedding_batch_identity is not None:
-        embedding_context["family_adapter"] = _fingerprint_jsonable(_embedding_batch_identity)
-        if allowed_unsupported_pooling:
-            embedding_context["family_adapter_pooling_override"] = sorted(
-                allowed_unsupported_pooling
-            )
-
-    tokenizer_metadata = _tokenizer_metadata(model, tokenizer)
-    (
-        input_fingerprint,
-        run_fingerprint,
-        resolved_model_state_fingerprint,
-        model_state_fingerprint_source,
-    ) = _run_fingerprint(
-        model,
-        records,
-        pooling=pooling_names,
-        full_embeddings=full_embeddings,
-        max_length=max_length,
-        truncate=truncate,
-        dtype=dtype,
-        model_kwargs=model_kwargs,
-        tokenizer_metadata=tokenizer_metadata,
-        model_state_fingerprint=model_state_fingerprint,
-        persist_output=output is not None,
-        embedding_context=embedding_context,
-        batch_size=batch_size,
-        batch_window_size=resolved_batch_window_size,
-        max_tokens_per_batch=max_tokens_per_batch,
-    )
-    output_already_exists = output is not None and _output_exists(output, format)
-    existing: EmbeddingResult | None = None
-    start_position = 0
-    if output is not None and resume and output_already_exists:
-        if format == "sqlite":
-            try:
-                existing = load_sqlite_result(output, run_id=run_fingerprint)
-            except KeyError:
-                existing = load_result(output, format=format)
-        else:
-            existing = load_result(output, format=format)
-        if existing.metadata.get("fingerprint_schema_version") != (_RUN_FINGERPRINT_SCHEMA_VERSION):
-            raise ValueError(
-                "Existing embeddings use an incompatible run fingerprint schema; "
-                "choose another output or set resume=False."
-            )
-        if existing.metadata.get("run_fingerprint") != run_fingerprint:
-            raise ValueError(
-                "Existing embeddings were produced by a different run fingerprint; "
-                "choose another output or set resume=False."
-            )
-        if len(existing) > len(records):
-            raise ValueError(
-                "Existing embeddings are not an ordered prefix of the requested inputs."
-            )
-        prefix_matches = all(
-            (observed.id, observed.sequence) == (expected.id, expected.sequence)
-            for expected, observed in zip(records, existing, strict=False)
-        )
-        if not prefix_matches:
-            raise ValueError(
-                "Existing embeddings are not an ordered prefix of the requested inputs."
-            )
-        if len(existing) == len(records) and existing.metadata.get("complete", True):
-            return existing
-        start_position = len(existing)
-
-    sqlite_run_id: str | None = None
-    sqlite_replace_on_first_commit = False
-    sqlite_initial_metadata: dict[str, Any] | None = None
-    if output is not None and format == "sqlite":
-        sqlite_initial_metadata = {
-            "format_version": 1,
-            "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-            "run_fingerprint": run_fingerprint,
-            "input_fingerprint": input_fingerprint,
-            "model_state_fingerprint": resolved_model_state_fingerprint,
-            "model_state_fingerprint_source": model_state_fingerprint_source,
-            "complete": False,
-        }
-        sqlite_run_id = run_fingerprint
-        if not resume and output_already_exists:
-            try:
-                load_sqlite_result(output, run_id=run_fingerprint)
-            except KeyError:
-                pass
-            else:
-                # Keep an exact prior run readable until replacement inference
-                # has produced the first complete commit window.
-                sqlite_replace_on_first_commit = True
-        if not sqlite_replace_on_first_commit:
-            initialize_sqlite_run(
-                output,
-                sqlite_initial_metadata,
-                resume=resume,
-            )
-
-    stream_safetensors = output is not None and format == "safetensors"
-    attention_backend = _attention_backend(model)
-    output_records: list[EmbeddingRecord] = (
-        [] if sqlite_run_id is not None or stream_safetensors else list(existing or ())
-    )
-    output_descriptors: list[dict[str, Any]] | None = [] if output is None else None
-    pool_slices: dict[str, tuple[int, int]] = {}
-    if existing and pooler is not None:
-        pooled_width = existing[0].load_tensor().shape[-1]
-        if pooled_width % len(pooling_names) != 0:
-            raise ValueError("Stored pooled width is inconsistent with pooling metadata.")
-        pool_slices = pooler.output_slices(pooled_width // len(pooling_names))
-
-    safetensors_writer: SafetensorsStreamWriter | None = None
-    if stream_safetensors:
-        if output is None:
-            raise RuntimeError("Safetensors streaming was enabled without an output destination.")
-        transactional_overwrite = output_already_exists and not resume
-        safetensors_writer = SafetensorsStreamWriter(
-            output,
-            {
-                "format_version": 1,
-                "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-                "run_fingerprint": run_fingerprint,
-                "input_fingerprint": input_fingerprint,
-                "model_state_fingerprint": resolved_model_state_fingerprint,
-                "model_state_fingerprint_source": model_state_fingerprint_source,
-                "complete": False,
-            },
-            shard_size=shard_size,
-            existing=existing or (),
-            reuse_existing=bool(resume and existing is not None),
-            publish_initial=not transactional_overwrite,
-            publish_incremental=not transactional_overwrite,
-        )
-    need_attentions = "parti" in pooling_names
-
-    config = getattr(model, "config", None)
-    model_type = str(getattr(config, "model_type", "")).lower()
-    resolved_tokenizer = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None)
-    with _temporary_eval(model), torch.inference_mode():
-        for window_start in range(start_position, len(records), resolved_batch_window_size):
-            window_stop = min(window_start + resolved_batch_window_size, len(records))
-            window_records = records[window_start:window_stop]
-            if not isinstance(window_records, Sequence):
-                raise RuntimeError("The immutable embedding spool returned a non-sequence window.")
-            window_results: dict[int, EmbeddingRecord] = {}
-            for local_positions in _planned_batches(
-                window_records,
-                range(len(window_records)),
-                batch_size=batch_size,
-                max_tokens_per_batch=max_tokens_per_batch,
-                max_length=max_length,
-                truncate=truncate,
-            ):
-                batch_positions = [window_start + position for position in local_positions]
-                batch_records = [window_records[position] for position in local_positions]
-                sequences = [
-                    record.sequence[:max_length]
-                    if truncate and max_length is not None
-                    else record.sequence
-                    for record in batch_records
-                ]
-                batch_model_kwargs = dict(model_kwargs)
-                if model_type == "fast_ankh" or hidden_state_source == "decoder":
-                    batch_model_kwargs["hidden_state_source"] = hidden_state_source
-                    if normalized_decoder_inputs is not None:
-                        batch_model_kwargs["decoder_inputs"] = [
-                            normalized_decoder_inputs[position] for position in batch_positions
-                        ]
-                    if decoder_input_ids is not None:
-                        # decoder_input_ids: (n_records, l_decoder)
-                        indices = torch.tensor(  # (b,)
-                            batch_positions,
-                            device=decoder_input_ids.device,
-                            dtype=torch.long,
-                        )
-                        batch_model_kwargs["decoder_input_ids"] = (  # (b, l_decoder)
-                            decoder_input_ids.index_select(0, indices)
-                        )
-                    if decoder_attention_mask is not None:
-                        # decoder_attention_mask: (n_records, l_decoder)
-                        indices = torch.tensor(  # (b,)
-                            batch_positions,
-                            device=decoder_attention_mask.device,
-                            dtype=torch.long,
-                        )
-                        batch_model_kwargs["decoder_attention_mask"] = (
-                            decoder_attention_mask.index_select(0, indices)  # (b, l_decoder)
-                        )
-                custom_batch = _embedding_batch_fn or getattr(model, "_embedding_batch", None)
-                if custom_batch is not None:
-                    if model_type == "fast_ankh":
-                        batch = custom_batch(
-                            sequences,
-                            tokenizer=resolved_tokenizer,
-                            max_length=max_length,
-                            truncate=truncate,
-                            need_attentions=need_attentions,
-                            **batch_model_kwargs,
-                        )
-                    else:
-                        batch = custom_batch(sequences, **batch_model_kwargs)
-                    if not isinstance(batch, EmbeddingBatch):
-                        raise TypeError("_embedding_batch must return EmbeddingBatch.")
-                else:
-                    batch = _generic_embedding_batch(
-                        model,
-                        sequences,
-                        tokenizer=tokenizer,
-                        max_length=max_length,
-                        truncate=truncate,
-                        need_attentions=need_attentions,
-                        model_kwargs=batch_model_kwargs,
-                    )
-                X = batch.X  # (b, l, d) or (b, n_states, l, d)
-                raw_mask = batch.residue_mask  # (b, l)
-                if not isinstance(X, Tensor) or not isinstance(raw_mask, Tensor):
-                    raise TypeError("Embedding batches must provide Tensor X and residue_mask.")
-                if X.is_meta or raw_mask.is_meta:
-                    raise ValueError("Embedding batches cannot contain meta tensors.")
-                if not X.is_floating_point():
-                    raise TypeError("Embedding batches must use a floating-point X dtype.")
-                if raw_mask.is_complex() or not bool(torch.isfinite(raw_mask).all()):
-                    raise ValueError("Embedding residue_mask must contain finite binary values.")
-                if not bool(((raw_mask == 0) | (raw_mask == 1)).all()):
-                    raise ValueError("Embedding residue_mask must contain finite binary values.")
-                M = raw_mask.to(device=X.device, dtype=torch.bool)  # (b, l)
-                valid_X_shape = (
-                    X.ndim == 3
-                    and X.shape[0] == len(batch_records)
-                    and X.shape[-1] > 0
-                    and M.shape == X.shape[:2]
-                )
-                valid_all_states_shape = (
-                    X.ndim == 4
-                    and store_all_hidden_states
-                    and full_embeddings
-                    and X.shape[0] == len(batch_records)
-                    and X.shape[1] > 0
-                    and X.shape[-1] > 0
-                    and M.shape == (X.shape[0], X.shape[2])
-                )
-                if not (valid_X_shape or valid_all_states_shape):
-                    raise ValueError(
-                        "Embedding batches must provide X with shape (b, l, d), or "
-                        "(b, states, l, d) when storing all hidden states, and "
-                        "residue_mask with shape (b, l)."
-                    )
-                if not bool(M.any(dim=1).all()):
-                    raise ValueError("Every embedding sample must contain a biological residue.")
-                finite_selected = (  # X.shape
-                    torch.isfinite(X) | ~M.unsqueeze(-1)
-                    if X.ndim == 3
-                    else torch.isfinite(X) | ~M[:, None, :, None]
-                )
-                if not bool(finite_selected.all()):
-                    raise ValueError("Biological residue embeddings produced non-finite output.")
-                if need_attentions:
-                    # Validate the biological graph only after mask integrity is established.
-                    _validate_parti_length(M)
-                if dtype is not None:
-                    X = X.to(dtype=dtype)  # unchanged shape
-
-                if full_embeddings:
-                    if X.ndim == 4:
-                        values = [
-                            X_i[:, M_i, :].detach().cpu()  # (n_states, r_i, d)
-                            for X_i, M_i in zip(X, M, strict=True)
-                        ]
-                    else:
-                        values = [
-                            X_i[M_i].detach().cpu()  # (r_i, d)
-                            for X_i, M_i in zip(X, M, strict=True)
-                        ]
-                else:
-                    if pooler is None:
-                        raise RuntimeError(
-                            "Pooled embedding output was requested without an initialized pooler."
-                        )
-                    Y = pooler(  # (b, n_poolers * d)
-                        X,
-                        M,
-                        attentions=batch.attentions,
-                        attention_backend=attention_backend,
-                    )
-                    pool_slices = pooler.output_slices(X.shape[-1])
-                    values = list(Y.detach().cpu().unbind(0))  # each: (n_poolers * d,)
-                for position, record, value in zip(
-                    batch_positions, batch_records, values, strict=True
-                ):
-                    window_results[position] = EmbeddingRecord(record.id, record.sequence, value)
-
-            new_records = [
-                window_results[position] for position in range(window_start, window_stop)
-            ]
-            if output_descriptors is not None:
-                output_descriptors.extend(
-                    _output_descriptor(window_start + offset, record)
-                    for offset, record in enumerate(new_records)
-                )
-            if output is not None and sqlite_run_id is not None:
-                append_sqlite_records(
-                    output,
-                    sqlite_run_id,
-                    window_start,
-                    new_records,
-                    replace_metadata=(
-                        sqlite_initial_metadata if sqlite_replace_on_first_commit else None
-                    ),
-                )
-                sqlite_replace_on_first_commit = False
-            elif safetensors_writer is not None:
-                safetensors_writer.append(new_records)
-            else:
-                output_records.extend(new_records)
-
-    software_versions = _software_versions()
-    projection = getattr(model, "embedding_projection", None)
-    resolved_layer = getattr(
-        model,
-        "embedding_layer",
-        model_kwargs.get("hidden_state_index", -1),
-    )
-    token_policy = getattr(
-        model,
-        "embedding_token_policy",
-        {
-            "unit": "residue",
-            "include": ["biological residues"],
-            "exclude": [
-                "BOS",
-                "EOS",
-                "padding",
-                "chain delimiters",
-                "non-protein tokens",
-            ],
-        },
-    )
-    model_identity = _model_identity_metadata(model)
-    metadata: dict[str, Any] = {
-        "format_version": 1,
-        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-        "run_fingerprint": run_fingerprint,
-        "input_fingerprint": input_fingerprint,
-        "model_state_fingerprint": resolved_model_state_fingerprint,
-        "model_state_fingerprint_source": model_state_fingerprint_source,
-        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
-        **model_identity,
-        "dtype": str(dtype).removeprefix("torch.") if dtype is not None else "model",
-        "attention_backend": attention_backend,
-        "attention_kernel": _attention_kernel_metadata(attention_backend),
-        "layer": resolved_layer,
-        "projection": projection,
-        "esmc_source": getattr(model, "_esmc_source", None),
-        "esmc_revision": getattr(model, "_esmc_source_revision", None),
-        "esmc_files": getattr(model, "_esmc_source_files", None),
-        "token_policy": token_policy,
-        "tokenizer": tokenizer_metadata,
-        **embedding_context,
-        "pooling": list(pooling_names),
-        "pool_slices": pool_slices,
-        "full_embeddings": full_embeddings,
-        "max_length": max_length,
-        "truncate": truncate,
-        "truncation": {"enabled": truncate, "max_length": max_length},
-        "batching": {
-            "batch_size": batch_size,
-            "batch_window_size": resolved_batch_window_size,
-            "max_tokens_per_batch": max_tokens_per_batch,
-            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
-            "ordering": "bounded-length-bucketed-stable-output",
-            "resume_commit_granularity": (
-                "not-applicable"
-                if output is None
-                else "batch-window"
-                if format == "sqlite"
-                else "shard-flush"
-            ),
-        },
-        "residue_mask_policy": "biological-residues-only",
-        "record_count": len(records),
-        "descriptor_index": (
-            "memory-metadata"
-            if output is None
-            else "sqlite-records"
-            if format == "sqlite"
-            else "safetensors-generation-index"
-        ),
-        "storage_format": format if output is not None else "memory",
-        "software": software_versions,
-        "execution": _execution_identity_metadata(model),
-        "adapter": _adapter_identity_metadata(model),
-        "torch_version": software_versions["torch"],
-        "transformers_version": software_versions["transformers"],
-        "complete": True,
-    }
-    if output_descriptors is not None:
-        metadata["outputs"] = output_descriptors
-        metadata["tensor_hashes"] = [item["sha256"] for item in output_descriptors]
-    status = getattr(model, "esmc_precision_status", None)
-    if status is not None:
-        metadata["esmc_precision"] = status.as_dict() if hasattr(status, "as_dict") else status
-    if output is not None and sqlite_run_id is not None:
-        update_sqlite_run_metadata(output, sqlite_run_id, metadata)
-        return load_sqlite_result(output, run_id=sqlite_run_id)
-    if safetensors_writer is not None:
-        return safetensors_writer.publish(complete=True, metadata=metadata)
-    result = EmbeddingResult(output_records, metadata)
-    if output is not None:
-        return save_result(result, output, format=format, shard_size=shard_size)
-    return result
-
-
-class EmbeddingMixin:
-    """Small delegation mixin shared by FastPLMs model classes."""
-
-    def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult:
-        return embed_dataset(self, inputs, **kwargs)
-
-
-__all__ = [
-    "EmbeddingMixin",
-    "embed_dataset",
-    "iter_fasta",
-    "parse_fasta",
-    "select_hidden_state_embeddings",
-]
+from . import identity
+from .batches import (
+    BatchExecutor,
+    _residue_embeddings as _residue_embeddings,
+    _temporary_eval,
+    select_hidden_state_embeddings as select_hidden_state_embeddings,
+)
+from .identity import (
+    _RUN_FINGERPRINT_SCHEMA_VERSION,
+    _adapter_identity_metadata,
+    _attention_backend,
+    _attention_kernel_metadata,
+    _embedding_context,
+    _execution_identity_metadata,
+    _fingerprint_jsonable,
+    _model_identity_metadata,
+    _run_fingerprint,
+    _tokenizer_metadata,
+)
+from .inputs import (
+    _InputSpool,
+    _normalize_inputs,
+    _validate_untruncated_lengths,
+    iter_fasta as iter_fasta,
+    parse_fasta as parse_fasta,
+)
+from .output import EmbeddingOutput
+from .pooling import Pooler
+from .types import EmbeddingBatch, EmbeddingInput, EmbeddingResult
+
+
+_DEFAULT_BATCH_WINDOW_MULTIPLIER = 16
+_SUPPORTED_STORAGE_FORMATS = frozenset({"safetensors", "sqlite"})
+
+
+def embed_dataset(
+    model: Any,
+    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
+    *,
+    batch_size: int = 2,
+    pooling: str | Sequence[str] | None = None,
+    full_embeddings: bool = False,
+    output: str | Path | None = None,
+    format: str = "safetensors",
+    resume: bool = True,
+    tokenizer: Any | None = None,
+    max_length: int | None = None,
+    truncate: bool = True,
+    dtype: torch.dtype | None = torch.float32,
+    shard_size: int = 2 * 1024**3,
+    model_state_fingerprint: str | None = None,
+    batch_window_size: int | None = None,
+    max_tokens_per_batch: int | None = None,
+    hidden_state_source: str = "encoder",
+    decoder_inputs: Sequence[str] | None = None,
+    decoder_input_ids: Tensor | None = None,
+    decoder_attention_mask: Tensor | None = None,
+    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None = None,
+    _embedding_batch_identity: Mapping[str, Any] | None = None,
+    _allowed_unsupported_pooling: Sequence[str] = (),
+    **model_kwargs: Any,
+) -> EmbeddingResult:
+    """Embed protein sequences with stable ordering and residue-only pooling."""
+
+    for name, value in (
+        ("batch_size", batch_size),
+        ("shard_size", shard_size),
+    ):
+        if not isinstance(value, int) or isinstance(value, bool):
+            raise TypeError(f"{name} must be a positive integer.")
+        if value <= 0:
+            raise ValueError(f"{name} must be a positive integer.")
+    for optional_name, optional_value in (
+        ("max_length", max_length),
+        ("max_tokens_per_batch", max_tokens_per_batch),
+        ("batch_window_size", batch_window_size),
+    ):
+        if optional_value is not None and (
+            not isinstance(optional_value, int) or isinstance(optional_value, bool)
+        ):
+            raise TypeError(f"{optional_name} must be a positive integer when provided.")
+        if optional_value is not None and optional_value <= 0:
+            raise ValueError(f"{optional_name} must be a positive integer when provided.")
+    for name, value in (
+        ("full_embeddings", full_embeddings),
+        ("resume", resume),
+        ("truncate", truncate),
+    ):
+        if not isinstance(value, bool):
+            raise TypeError(f"{name} must be a boolean.")
+    if not isinstance(format, str):
+        raise TypeError("format must be a string.")
+    if output is not None and not isinstance(output, (str, Path)):
+        raise TypeError("output must be a path or None.")
+    if model_state_fingerprint is not None and (
+        not isinstance(model_state_fingerprint, str) or not model_state_fingerprint
+    ):
+        raise ValueError("model_state_fingerprint must be a non-empty string when provided.")
+    if hidden_state_source not in {"encoder", "decoder"}:
+        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
+    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
+    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
+        raise TypeError("hidden_state_index must be an integer.")
+    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
+    if not isinstance(store_all_hidden_states, bool):
+        raise TypeError("store_all_hidden_states must be a boolean.")
+    if decoder_input_ids is not None:
+        if not isinstance(decoder_input_ids, Tensor):
+            raise TypeError("decoder_input_ids must be a tensor.")
+        if decoder_input_ids.is_meta:
+            raise ValueError("decoder_input_ids cannot be a meta tensor.")
+        if decoder_input_ids.ndim != 2 or decoder_input_ids.shape[1] == 0:
+            raise ValueError("decoder_input_ids must have non-empty shape (batch, sequence).")
+        if decoder_input_ids.dtype not in {torch.int32, torch.int64}:
+            raise TypeError("decoder_input_ids must use torch.int32 or torch.int64.")
+    if decoder_attention_mask is not None:
+        if not isinstance(decoder_attention_mask, Tensor):
+            raise TypeError("decoder_attention_mask must be a tensor.")
+        if decoder_attention_mask.is_meta:
+            raise ValueError("decoder_attention_mask cannot be a meta tensor.")
+        if decoder_attention_mask.is_complex() or not bool(
+            torch.isfinite(decoder_attention_mask).all()
+        ):
+            raise ValueError("decoder_attention_mask must contain finite binary values.")
+        if not bool(((decoder_attention_mask == 0) | (decoder_attention_mask == 1)).all()):
+            raise ValueError("decoder_attention_mask must contain finite binary values.")
+    pooling_names = (
+        (("mean",) if not full_embeddings else ())
+        if pooling is None
+        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
+    )
+    if full_embeddings and pooling is not None:
+        raise ValueError("full_embeddings=True cannot be combined with pooling.")
+    if not full_embeddings and not pooling_names:
+        raise ValueError("pooling is required unless full_embeddings=True.")
+    pooler = Pooler(pooling_names) if pooling_names else None
+
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive.")
+    if format == "pth" or (output is not None and Path(output).suffix.lower() == ".pth"):
+        raise ValueError("Writing pickle-based .pth embeddings is not supported.")
+    if format not in _SUPPORTED_STORAGE_FORMATS:
+        raise ValueError("format must be 'safetensors' or 'sqlite'.")
+    if max_length is not None and max_length <= 0:
+        raise ValueError("max_length must be positive when provided.")
+    if max_tokens_per_batch is not None and max_tokens_per_batch <= 0:
+        raise ValueError("max_tokens_per_batch must be positive when provided.")
+    if not isinstance(dtype, (torch.dtype, type(None))):
+        raise TypeError("dtype must be a torch.dtype or None.")
+    if batch_window_size is not None and batch_window_size <= 0:
+        raise ValueError("batch_window_size must be positive when provided.")
+    if _embedding_batch_fn is not None and not callable(_embedding_batch_fn):
+        raise TypeError("_embedding_batch_fn must be callable when provided.")
+    if _embedding_batch_fn is not None and _embedding_batch_identity is None:
+        raise ValueError(
+            "_embedding_batch_identity is required with _embedding_batch_fn so persisted "
+            "runs bind the family-specific embedding behavior."
+        )
+    if _embedding_batch_identity is not None and not isinstance(_embedding_batch_identity, Mapping):
+        raise TypeError("_embedding_batch_identity must be a mapping when provided.")
+    if isinstance(_allowed_unsupported_pooling, (str, bytes)) or not isinstance(
+        _allowed_unsupported_pooling, Sequence
+    ):
+        raise TypeError("_allowed_unsupported_pooling must be a sequence of pooler names.")
+    if not all(isinstance(name, str) for name in _allowed_unsupported_pooling):
+        raise TypeError("_allowed_unsupported_pooling must contain only strings.")
+    allowed_unsupported_pooling = frozenset(_allowed_unsupported_pooling)
+    if allowed_unsupported_pooling and _embedding_batch_fn is None:
+        raise ValueError(
+            "_allowed_unsupported_pooling is only valid with a family-specific _embedding_batch_fn."
+        )
+    resolved_batch_window_size = (
+        batch_size * _DEFAULT_BATCH_WINDOW_MULTIPLIER
+        if batch_window_size is None
+        else batch_window_size
+    )
+    if resolved_batch_window_size < batch_size:
+        raise ValueError("batch_window_size must be at least batch_size.")
+    records = _normalize_inputs(inputs, disk_backed=output is not None)
+    _validate_untruncated_lengths(
+        records,
+        max_length=max_length,
+        truncate=truncate,
+    )
+    pooling_names = (
+        (("mean",) if not full_embeddings else ())
+        if pooling is None
+        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
+    )
+    if full_embeddings:
+        if pooling is not None:
+            raise ValueError("full_embeddings=True cannot be combined with pooling.")
+    elif not pooling_names:
+        raise ValueError("pooling is required unless full_embeddings=True.")
+    store_all_hidden_states = bool(model_kwargs.get("store_all_hidden_states", False))
+    if store_all_hidden_states and not full_embeddings:
+        raise ValueError("store_all_hidden_states=True requires full_embeddings=True.")
+
+    unsupported = set(getattr(model, "embedding_unsupported_pooling", ()))
+    unknown_pooling_overrides = allowed_unsupported_pooling.difference(unsupported)
+    if unknown_pooling_overrides:
+        raise ValueError(
+            "_allowed_unsupported_pooling may only override poolers declared unsupported "
+            f"by the model; unknown overrides: {sorted(unknown_pooling_overrides)}."
+        )
+    unsupported.difference_update(allowed_unsupported_pooling)
+    requested_unsupported = unsupported.intersection(pooling_names)
+    if requested_unsupported:
+        raise ValueError(
+            f"{model.__class__.__name__} does not support pooling operations "
+            f"{sorted(requested_unsupported)}."
+        )
+
+    # Constructing the pooler validates names and duplicate operations before
+    # any checkpoint hashing, tokenization, or inference occurs.
+    pooler = Pooler(pooling_names) if pooling_names else None
+    embedding_context, normalized_decoder_inputs = _embedding_context(
+        model,
+        records,
+        hidden_state_source=hidden_state_source,
+        decoder_inputs=decoder_inputs,
+        decoder_input_ids=decoder_input_ids,
+        decoder_attention_mask=decoder_attention_mask,
+        model_kwargs=model_kwargs,
+    )
+    if _embedding_batch_identity is not None:
+        embedding_context["family_adapter"] = _fingerprint_jsonable(_embedding_batch_identity)
+        if allowed_unsupported_pooling:
+            embedding_context["family_adapter_pooling_override"] = sorted(
+                allowed_unsupported_pooling
+            )
+
+    # A pending automatic attention request settles here, inside the caller's
+    # autocast context, so the fingerprint records the backend that executes.
+    attention_resolution = getattr(model, "attention_resolution", None)
+    if attention_resolution is not None and attention_resolution.deferred:
+        model.resolve_attn_implementation()
+
+    tokenizer_metadata = _tokenizer_metadata(model, tokenizer)
+    (
+        input_fingerprint,
+        run_fingerprint,
+        resolved_model_state_fingerprint,
+        model_state_fingerprint_source,
+    ) = _run_fingerprint(
+        model,
+        records,
+        pooling=pooling_names,
+        full_embeddings=full_embeddings,
+        max_length=max_length,
+        truncate=truncate,
+        dtype=dtype,
+        model_kwargs=model_kwargs,
+        tokenizer_metadata=tokenizer_metadata,
+        model_state_fingerprint=model_state_fingerprint,
+        persist_output=output is not None,
+        embedding_context=embedding_context,
+        batch_size=batch_size,
+        batch_window_size=resolved_batch_window_size,
+        max_tokens_per_batch=max_tokens_per_batch,
+    )
+    destination = EmbeddingOutput(
+        records,
+        output=output,
+        format=format,
+        resume=resume,
+        shard_size=shard_size,
+        run_fingerprint=run_fingerprint,
+        input_fingerprint=input_fingerprint,
+        model_state_fingerprint=resolved_model_state_fingerprint,
+        model_state_fingerprint_source=model_state_fingerprint_source,
+        pooler=pooler,
+        pooling_names=pooling_names,
+    )
+    if destination.completed is not None:
+        return destination.completed
+
+    attention_backend = _attention_backend(model)
+    executor = BatchExecutor(
+        model=model,
+        batch_size=batch_size,
+        max_tokens_per_batch=max_tokens_per_batch,
+        max_length=max_length,
+        truncate=truncate,
+        model_kwargs=model_kwargs,
+        hidden_state_source=hidden_state_source,
+        normalized_decoder_inputs=normalized_decoder_inputs,
+        decoder_input_ids=decoder_input_ids,
+        decoder_attention_mask=decoder_attention_mask,
+        _embedding_batch_fn=_embedding_batch_fn,
+        tokenizer=tokenizer,
+        store_all_hidden_states=store_all_hidden_states,
+        full_embeddings=full_embeddings,
+        dtype=dtype,
+        pooler=pooler,
+        attention_backend=attention_backend,
+        need_attentions="parti" in pooling_names,
+    )
+    pool_slices = destination.pool_slices
+    with _temporary_eval(model), torch.inference_mode():
+        for window_start in range(
+            destination.start_position, len(records), resolved_batch_window_size
+        ):
+            window_stop = min(window_start + resolved_batch_window_size, len(records))
+            window_records = records[window_start:window_stop]
+            if not isinstance(window_records, Sequence):
+                raise RuntimeError("The immutable embedding spool returned a non-sequence window.")
+            new_records, pool_slices = executor.run_window(
+                window_records, window_start=window_start
+            )
+            destination.append(window_start, new_records)
+
+    software_versions = identity._software_versions()
+    projection = getattr(model, "embedding_projection", None)
+    resolved_layer = getattr(
+        model,
+        "embedding_layer",
+        model_kwargs.get("hidden_state_index", -1),
+    )
+    token_policy = getattr(
+        model,
+        "embedding_token_policy",
+        {
+            "unit": "residue",
+            "include": ["biological residues"],
+            "exclude": [
+                "BOS",
+                "EOS",
+                "padding",
+                "chain delimiters",
+                "non-protein tokens",
+            ],
+        },
+    )
+    model_identity = _model_identity_metadata(model)
+    metadata: dict[str, Any] = {
+        "format_version": 1,
+        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+        "run_fingerprint": run_fingerprint,
+        "input_fingerprint": input_fingerprint,
+        "model_state_fingerprint": resolved_model_state_fingerprint,
+        "model_state_fingerprint_source": model_state_fingerprint_source,
+        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
+        **model_identity,
+        "dtype": str(dtype).removeprefix("torch.") if dtype is not None else "model",
+        "attention_backend": attention_backend,
+        "attention_kernel": _attention_kernel_metadata(attention_backend),
+        "layer": resolved_layer,
+        "projection": projection,
+        "esmc_source": getattr(model, "_esmc_source", None),
+        "esmc_revision": getattr(model, "_esmc_source_revision", None),
+        "esmc_files": getattr(model, "_esmc_source_files", None),
+        "token_policy": token_policy,
+        "tokenizer": tokenizer_metadata,
+        **embedding_context,
+        "pooling": list(pooling_names),
+        "pool_slices": pool_slices,
+        "full_embeddings": full_embeddings,
+        "max_length": max_length,
+        "truncate": truncate,
+        "truncation": {"enabled": truncate, "max_length": max_length},
+        "batching": {
+            "batch_size": batch_size,
+            "batch_window_size": resolved_batch_window_size,
+            "max_tokens_per_batch": max_tokens_per_batch,
+            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
+            "ordering": "bounded-length-bucketed-stable-output",
+            "resume_commit_granularity": (
+                "not-applicable"
+                if output is None
+                else "batch-window"
+                if format == "sqlite"
+                else "shard-flush"
+            ),
+        },
+        "residue_mask_policy": "biological-residues-only",
+        "record_count": len(records),
+        "descriptor_index": (
+            "memory-metadata"
+            if output is None
+            else "sqlite-records"
+            if format == "sqlite"
+            else "safetensors-generation-index"
+        ),
+        "storage_format": format if output is not None else "memory",
+        "software": software_versions,
+        "execution": _execution_identity_metadata(model),
+        "adapter": _adapter_identity_metadata(model),
+        "torch_version": software_versions["torch"],
+        "transformers_version": software_versions["transformers"],
+        "complete": True,
+    }
+    if destination.output_descriptors is not None:
+        metadata["outputs"] = destination.output_descriptors
+        metadata["tensor_hashes"] = [item["sha256"] for item in destination.output_descriptors]
+    status = getattr(model, "esmc_precision_status", None)
+    if status is not None:
+        metadata["esmc_precision"] = status.as_dict() if hasattr(status, "as_dict") else status
+    return destination.finish(metadata)
+
+
+class EmbeddingMixin:
+    """Small delegation mixin shared by FastPLMs model classes."""
+
+    def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult:
+        return embed_dataset(self, inputs, **kwargs)
+
+
+__all__ = [
+    "EmbeddingMixin",
+    "embed_dataset",
+    "iter_fasta",
+    "parse_fasta",
+    "select_hidden_state_embeddings",
+]

fastplms/models.toml

@@ -11,6 +11,7 @@ revision = "db6b51744f0cd7061386442c09df890fc6d9f47e"
 version = 2
 expected_variant = "flash_attn2"
 dtypes = ["bfloat16"]
+min_cuda_capability = [8, 0]
 
 [[attention_kernels]]
 implementation = "flash_attention_3"
@@ -19,6 +20,7 @@ revision = "43f0bd269777115d94ff826e0d113ce9c1c9087b"
 version = 1
 expected_variant = "flash_attn3"
 dtypes = ["bfloat16"]
+min_cuda_capability = [8, 0]
 
 [[runtime_assets]]
 id = "esmfold2_ccd"
@@ -161,6 +163,8 @@ extra = "core"
 reference_container = "reference-esm2"
 reference_adapter = "tests.parity.support.reference_adapters.esm2"
 attention = ["eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"]
+attention_auto_order = ["sdpa"]
+attention_auto_evidence = "docs/evidence/attention/backend_latency.json"
 dtypes = ["float32", "bfloat16"]
 bf16_execution = "fp32_parameters_autocast"
 precisions = ["default"]
@@ -185,6 +189,8 @@ extra = "core"
 reference_container = "reference-biohub-esm"
 reference_adapter = "tests.parity.support.reference_adapters.esm_plusplus"
 attention = ["eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"]
+attention_auto_order = ["sdpa"]
+attention_auto_evidence = "docs/evidence/attention/backend_latency.json"
 dtypes = ["float32", "bfloat16"]
 bf16_execution = "static_parameters"
 precisions = ["default", "fp8"]
@@ -210,6 +216,7 @@ extra = "core"
 reference_container = "reference-biohub-esm"
 reference_adapter = "tests.parity.support.reference_adapters.esm3"
 attention = ["eager", "sdpa", "flex_attention"]
+attention_auto_order = ["sdpa"]
 dtypes = ["float32", "bfloat16"]
 bf16_execution = "fp32_parameters_autocast"
 precisions = ["default"]
@@ -234,6 +241,7 @@ extra = "core"
 reference_container = "reference-e1"
 reference_adapter = "tests.parity.support.reference_adapters.e1"
 attention = ["sdpa", "flex_attention"]
+attention_auto_order = ["sdpa"]
 dtypes = ["float32", "bfloat16"]
 bf16_execution = "static_parameters"
 precisions = ["default"]
@@ -260,6 +268,8 @@ extra = "core"
 reference_container = "reference-dplm"
 reference_adapter = "tests.parity.support.reference_adapters.dplm"
 attention = ["eager", "sdpa", "flex_attention", "flash_attention_3"]
+attention_auto_order = ["sdpa"]
+attention_auto_evidence = "docs/evidence/attention/backend_latency.json"
 dtypes = ["float32", "bfloat16"]
 bf16_execution = "fp32_parameters_autocast"
 precisions = ["default"]
@@ -284,6 +294,7 @@ extra = "core"
 reference_container = "reference-dplm"
 reference_adapter = "tests.parity.support.reference_adapters.dplm2"
 attention = ["sdpa"]
+attention_auto_order = ["sdpa"]
 dtypes = ["float32", "bfloat16"]
 bf16_execution = "fp32_parameters_autocast"
 precisions = ["default"]
@@ -309,6 +320,7 @@ extra = "core"
 reference_container = "reference-ankh"
 reference_adapter = "tests.parity.support.reference_adapters.ankh"
 attention = ["eager", "sdpa"]
+attention_auto_order = ["sdpa", "eager"]
 dtypes = ["float32", "bfloat16"]
 bf16_execution = "static_parameters"
 precisions = ["default"]
@@ -358,6 +370,7 @@ extra = "structure"
 reference_container = "reference-esmfold"
 reference_adapter = "tests.parity.support.reference_adapters.esmfold"
 attention = ["eager", "sdpa", "flex_attention"]
+attention_auto_order = ["sdpa"]
 dtypes = ["float32", "bfloat16"]
 bf16_execution = "fp32_parameters_autocast"
 precisions = ["default"]
@@ -383,6 +396,7 @@ extra = "structure"
 reference_container = "reference-esmfold2"
 reference_adapter = "tests.parity.support.reference_adapters.esmfold2"
 attention = ["eager", "sdpa", "flex_attention"]
+attention_auto_order = ["sdpa"]
 dtypes = ["float32", "bfloat16"]
 bf16_execution = "fp32_parameters_autocast"
 precisions = ["auto", "fp32", "bf16", "fp8"]

fastplms/models/_esm_rotary.py

@@ -56,13 +56,30 @@ class RotaryEmbedding(nn.Module):
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # tensor: (..., l, d)
         seq_len = tensor.shape[seq_dimension]
-        cache_stale = (
-            self._cos_cached is None
-            or self._sin_cached is None
-            or self._seq_len_cached != seq_len
-            or self._cos_cached.device != tensor.device
+        # Row t depends only on t and ``inv_freq``, so the first l rows of a
+        # longer table equal a table built at length l. A shorter request can
+        # therefore reuse the cache, provided ``inv_freq`` still has the dtype
+        # the cache was built from. A length change used to rebuild from the
+        # current ``inv_freq``; the dtype test preserves that after ``.to(dtype)``.
+        # Autograd must save these factors, so inference-created tables cannot
+        # be reused in a gradient-enabled forward, even when the module is in eval mode.
+        cached_prefix_is_valid = (
+            self._cos_cached is not None
+            and self._sin_cached is not None
+            and self._seq_len_cached is not None
+            and self._cos_cached.device == tensor.device
+            and not (
+                torch.is_grad_enabled()
+                and (self._cos_cached.is_inference() or self._sin_cached.is_inference())
+            )
+            and (
+                seq_len == self._seq_len_cached
+                or (
+                    seq_len < self._seq_len_cached and self._cos_cached.dtype == self.inv_freq.dtype
+                )
+            )
         )
-        if cache_stale:
+        if not cached_prefix_is_valid:
             self._seq_len_cached = seq_len
             positions = torch.arange(seq_len, device=tensor.device).type_as(  # (l,)
                 self.inv_freq
@@ -74,7 +91,8 @@ class RotaryEmbedding(nn.Module):
 
         assert self._cos_cached is not None
         assert self._sin_cached is not None
-        return self._cos_cached, self._sin_cached  # (1, 1, l, d), (1, 1, l, d)
+        # Callers still receive exactly l rows: (1, 1, l, d), (1, 1, l, d)
+        return self._cos_cached[:, :, :seq_len, :], self._sin_cached[:, :, :seq_len, :]
 
     def forward(
         self,

fastplms/models/esm_plusplus/modeling_esm_plusplus.py

@@ -35,10 +35,12 @@ try:
         AttentionBackend,
         BlockMask,
         FastPLMsAttentionMixin,
+        FlashPaddingLayout,
         _get_flex_attention_fn,
         _get_flex_block_mask,
         flex_attention,
         get_attention_mask,
+        get_flash_padding_layout,
         kernels_flash_attention_func,
         resolve_attention_backend,
         resolve_attention_backend_for_call,
@@ -52,11 +54,13 @@ except ModuleNotFoundError as error:
         "EmbeddingMixin",
         "FastPLMsAttentionMixin",
         "FastPLMTestTimeTrainingMixin",
+        "FlashPaddingLayout",
         "Pooler",
         "_get_flex_attention_fn",
         "_get_flex_block_mask",
         "flex_attention",
         "get_attention_mask",
+        "get_flash_padding_layout",
         "kernels_flash_attention_func",
         "resolve_attention_backend",
         "resolve_attention_backend_for_call",
@@ -296,13 +300,27 @@ def apply_rotary_emb_torch(
         raise AssertionError("rotary width exceeds the attention head dimension")
 
     token_count = x.shape[1]
-    cos_full = torch.cat((cos[:token_count], cos[:token_count]), dim=-1).unsqueeze(1)
-    sin_full = torch.cat((sin[:token_count], sin[:token_count]), dim=-1).unsqueeze(1)
-    x_rotary = x[..., :rotary_width]
-    y_rotary = x_rotary * cos_full + rotate_half(x_rotary, interleaved) * sin_full
+    cos_full = torch.cat((cos[:token_count], cos[:token_count]), dim=-1)  # (l, d_r)
+    sin_full = torch.cat((sin[:token_count], sin[:token_count]), dim=-1)  # (l, d_r)
+    return _rotate_with_full_tables(x, cos_full, sin_full, interleaved)
+
+
+def _rotate_with_full_tables(
+    x: torch.Tensor,
+    cos_full: torch.Tensor,
+    sin_full: torch.Tensor,
+    interleaved: bool,
+) -> torch.Tensor:
+    # x: (b, l, h, d); cos_full, sin_full: (l, d_r) with rotary width d_r <= d
+    rotary_width = cos_full.shape[-1]
+    x_rotary = x[..., :rotary_width]  # (b, l, h, d_r)
+    y_rotary = (  # (b, l, h, d_r)
+        x_rotary * cos_full.unsqueeze(1)
+        + rotate_half(x_rotary, interleaved) * sin_full.unsqueeze(1)
+    )
     if rotary_width == x.shape[-1]:
         return y_rotary
-    return torch.cat((y_rotary, x[..., rotary_width:]), dim=-1)
+    return torch.cat((y_rotary, x[..., rotary_width:]), dim=-1)  # (b, l, h, d)
 
 
 class RotaryEmbedding(torch.nn.Module):
@@ -344,6 +362,10 @@ class RotaryEmbedding(torch.nn.Module):
         self._sin_cached: torch.Tensor | None = None
         self._cos_k_cached: torch.Tensor | None = None
         self._sin_k_cached: torch.Tensor | None = None
+        # Both halves of a head rotate by the same angles. Q and K of every layer
+        # read these full-width tables instead of concatenating the halves again.
+        self._cos_full_cached: torch.Tensor | None = None
+        self._sin_full_cached: torch.Tensor | None = None
 
     def reset_parameters(self, device: torch.device | str | None = None) -> None:
         """Rebuild the non-persistent frequency buffers on ``device``."""
@@ -426,6 +448,8 @@ class RotaryEmbedding(torch.nn.Module):
         if self.scale is None:
             self._cos_cached = cos_angles.to(dtype)
             self._sin_cached = sin_angles.to(dtype)
+            self._cos_full_cached = torch.cat((self._cos_cached, self._cos_cached), dim=-1)
+            self._sin_full_cached = torch.cat((self._sin_cached, self._sin_cached), dim=-1)
             return
 
         centered_positions = (
@@ -459,11 +483,17 @@ class RotaryEmbedding(torch.nn.Module):
         if self.scale is not None:
             raise AssertionError("Scaled rotary embeddings are unsupported for ESMC.")
 
-        cos_angles = self._cos_cached
-        sin_angles = self._sin_cached
+        if self._cos_full_cached is None or self._sin_full_cached is None:
+            raise RuntimeError("Rotary cache initialization did not produce full-width tables.")
+        if 2 * self._cos_cached.shape[-1] > q.shape[-1]:
+            raise AssertionError("rotary width exceeds the attention head dimension")
+
+        token_count = q.shape[1]
+        cos_full = self._cos_full_cached[:token_count]  # (l, d_r)
+        sin_full = self._sin_full_cached[:token_count]  # (l, d_r)
         return (
-            apply_rotary_emb_torch(q, cos_angles, sin_angles, self.interleaved, True),
-            apply_rotary_emb_torch(k, cos_angles, sin_angles, self.interleaved, True),
+            _rotate_with_full_tables(q, cos_full, sin_full, self.interleaved),
+            _rotate_with_full_tables(k, cos_full, sin_full, self.interleaved),
         )
 
 
@@ -540,6 +570,7 @@ class MultiHeadAttention(nn.Module):
         flex_block_mask: BlockMask | None = None,
         output_attentions: bool = False,
         output_s_max: bool = False,
+        flash_padding_layout: FlashPaddingLayout | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
         # x: (b, l, d)
         qkv = self.layernorm_qkv(x)  # (b, l, 3 * d)
@@ -562,6 +593,7 @@ class MultiHeadAttention(nn.Module):
             flex_block_mask=flex_block_mask,
             output_attentions=output_attentions,
             output_s_max=output_s_max,
+            flash_padding_layout=flash_padding_layout,
         )
 
         output = self.out_proj(attn_output)
@@ -577,6 +609,7 @@ class MultiHeadAttention(nn.Module):
         flex_block_mask: BlockMask | None = None,
         output_attentions: bool = False,
         output_s_max: bool = False,
+        flash_padding_layout: FlashPaddingLayout | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
         if output_attentions:
             return self._manual_attn(
@@ -590,7 +623,7 @@ class MultiHeadAttention(nn.Module):
             return attn_output, None, s_max
         if self.attn_backend.is_flash:
             attn_output, attn_weights = self._kernels_flash_attn(
-                query_heads, key_heads, value_heads, attention_mask_2d
+                query_heads, key_heads, value_heads, attention_mask_2d, flash_padding_layout
             )
         elif self.attn_backend == AttentionBackend.FLEX:
             attn_output, attn_weights = self._flex_attn(
@@ -647,6 +680,7 @@ class MultiHeadAttention(nn.Module):
         key_heads: torch.Tensor,
         value_heads: torch.Tensor,
         attention_mask_2d: torch.Tensor | None = None,
+        flash_padding_layout: FlashPaddingLayout | None = None,
     ) -> tuple[torch.Tensor, None]:
         query_tokens = query_heads.transpose(1, 2).contiguous()
         key_tokens = key_heads.transpose(1, 2).contiguous()
@@ -658,6 +692,7 @@ class MultiHeadAttention(nn.Module):
             attention_mask_2d=attention_mask_2d,
             causal=False,
             implementation=self.attn_backend.value,
+            padding_layout=flash_padding_layout,
         )
         return rearrange(attn_output, "b s h d -> b s (h d)"), None
 
@@ -747,6 +782,7 @@ class UnifiedTransformerBlock(nn.Module):
         flex_block_mask: BlockMask | None = None,
         output_attentions: bool = False,
         output_s_max: bool = False,
+        flash_padding_layout: FlashPaddingLayout | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
         attn_output, attn_weights, s_max = self.attn(
             x,
@@ -755,6 +791,7 @@ class UnifiedTransformerBlock(nn.Module):
             flex_block_mask=flex_block_mask,
             output_attentions=output_attentions,
             output_s_max=output_s_max,
+            flash_padding_layout=flash_padding_layout,
         )
         x = x + self.dropout(attn_output) / self.scaling_factor
         x = x + self.dropout(self.ffn(x)) / self.scaling_factor
@@ -857,16 +894,20 @@ class TransformerStack(nn.Module):
         # Match the pinned Biohub Transformers contract: a supplied sequence_id
         # is authoritative and must encode padding as -1.  attention_mask is
         # ignored in that mode rather than intersected with the chain mask.
-        attention_mask_2d, attention_mask_4d, flex_block_mask = (
-            self._prepare_attention_masks(
-                attention_mask=attention_mask,
-                sequence_id=sequence_id,
-                batch_size=x.shape[0],
-                seq_len=x.shape[1],
-                device=x.device,
-                dtype=x.dtype,
-                output_attentions=bool(output_attentions),
-            )
+        attention_mask_2d, attention_mask_4d, flex_block_mask = self._prepare_attention_masks(
+            attention_mask=attention_mask,
+            sequence_id=sequence_id,
+            batch_size=x.shape[0],
+            seq_len=x.shape[1],
+            device=x.device,
+            dtype=x.dtype,
+            output_attentions=bool(output_attentions),
+        )
+        # A call that returns attention weights runs eager attention and needs no layout.
+        flash_padding_layout = (
+            None
+            if output_attentions
+            else get_flash_padding_layout(self.attention_backend, attention_mask_2d)
         )
 
         for layer_index, block in enumerate(self.blocks):
@@ -890,6 +931,7 @@ class TransformerStack(nn.Module):
                     flex_block_mask=flex_block_mask,
                     output_attentions=output_attentions,
                     output_s_max=output_s_max,
+                    flash_padding_layout=flash_padding_layout,
                 )
             else:
                 x, attn_weights, s_max = block(
@@ -899,6 +941,7 @@ class TransformerStack(nn.Module):
                     flex_block_mask=flex_block_mask,
                     output_attentions=output_attentions,
                     output_s_max=output_s_max,
+                    flash_padding_layout=flash_padding_layout,
                 )
 
             if attentions is not None:
@@ -1048,6 +1091,7 @@ class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel):
         "flash_attention_2",
         "flash_attention_3",
     )
+    _fastplms_attention_auto_order = ("sdpa",)
 
     def __init__(self, config: ESMplusplusConfig, *args: object, **kwargs: object) -> None:
         super().__init__(config, *args, **kwargs)

fastplms/models/esmfold2/attention.py

@@ -20,6 +20,7 @@ class ESMFold2AttentionMixin(FastPLMsAttentionMixin):
         "sdpa",
         "flex_attention",
     )
+    _fastplms_attention_auto_order = ("sdpa",)
 
     def __init__(self, config, *args, **kwargs) -> None:
         super().__init__(config, *args, **kwargs)

fastplms/models/esmfold2/configuration_esmfold2.py

@@ -1,336 +1,336 @@
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Configuration schema for release and experimental ESMFold2 checkpoints."""
-
-from __future__ import annotations
-
-from dataclasses import asdict, dataclass, field
-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"})
-
-
-def _esmc_backbone_checkpoint_ids() -> tuple[str, str]:
-    """Return the manifest-pinned official and FastPLMs ESMC repositories."""
-
-    from fastplms.registry import RegistryError, get_model_registry
-
-    registry = get_model_registry()
-    family = registry.families["esmfold2"]
-    if family.backbone_model is None:
-        raise RegistryError("families.esmfold2 must declare backbone_model.")
-    backbone = registry[family.backbone_model]
-    return backbone.official.repo_id, backbone.fast.repo_id
-
-
-def normalize_esmc_id(esmc_id: str) -> str:
-    """Resolve an official ESMC identifier to its FastPLMs checkpoint mirror."""
-
-    from fastplms.registry import get_model_registry
-
-    official_repo, fast_repo = _esmc_backbone_checkpoint_ids()
-    if esmc_id == official_repo:
-        return fast_repo
-    registry = get_model_registry()
-    for spec in registry.by_family("esmfold2"):
-        if spec.backbone is not None and spec.backbone_model is not None:
-            backbone = registry[spec.backbone_model]
-            if esmc_id in {spec.backbone.repo_id, backbone.official.repo_id}:
-                return backbone.fast.repo_id
-    return esmc_id
-
-
-def normalize_esmc_attention_implementation(
-    implementation: str | dict[str, str] | None,
-) -> str | None:
-    """Validate the ESMC backend and translate the historical ``flex`` name."""
-
-    if isinstance(implementation, dict):
-        if tuple(implementation) != ("",):
-            raise ValueError(
-                "ESMFold2 has one ESMC attention backbone; use a string or {'': implementation}."
-            )
-        implementation = 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(
-            f"Unsupported ESMFold2 attention implementation {canonical!r}; "
-            f"expected one of {expected}."
-        )
-    return canonical
-
-
-NestedConfig = TypeVar("NestedConfig")
-
-
-def _nested_config(value: Any, config_type: type[NestedConfig]) -> NestedConfig:
-    if isinstance(value, config_type):
-        return value
-    return config_type(**value) if isinstance(value, dict) else config_type()
-
-
-def _coerce_nested_field(
-    value: NestedConfig | dict[str, Any], config_type: type[NestedConfig]
-) -> NestedConfig:
-    """Convert serialized nested dictionaries while retaining supplied objects."""
-
-    return config_type(**value) if isinstance(value, dict) else value
-
-
-@dataclass
-class AtomAttentionConfig:
-    """Sliding-window atom attention and three-dimensional RoPE settings."""
-
-    d_atom: int = field(default=128)
-    d_token: int = field(default=768)
-    n_blocks: int = field(default=3)
-    n_heads: int = field(default=4)
-    swa_window_size: int = field(default=128)
-    expansion_ratio: int = field(default=2)
-    spatial_rope_base_frequency: float = field(default=20.0)
-    n_spatial_rope_pairs_per_axis: int = field(default=2)
-    n_uid_rope_pairs: int = field(default=10)
-    uid_rope_base_frequency: float = field(default=10000.0)
-
-
-@dataclass
-class DiffusionModuleConfig:
-    """Dimensions and depth of the coordinate diffusion network."""
-
-    sigma_data: float = field(default=16.0)
-    c_atom: int = field(default=128)
-    c_token: int = field(default=768)
-    c_z: int = field(default=256)
-    c_s_inputs: int = field(default=451)
-    fourier_dim: int = field(default=256)
-    relpos_r_max: int = field(default=32)
-    relpos_s_max: int = field(default=2)
-    atom_num_blocks: int = field(default=3)
-    atom_num_heads: int = field(default=4)
-    token_num_blocks: int = field(default=12)
-    token_num_heads: int = field(default=16)
-    transition_multiplier: int = field(default=2)
-
-
-@dataclass
-class FoldingTrunkConfig:
-    """Iterative pair/single trunk dimensions."""
-
-    n_layers: int = field(default=24)
-    n_heads: int = field(default=8)
-    dropout: float = field(default=0.0)
-
-
-@dataclass
-class InputsEmbedderConfig:
-    """Input feature width and atom encoder settings."""
-
-    d_inputs: int = field(default=451)
-    atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig)
-
-    def __post_init__(self) -> None:
-        self.atom_encoder = _coerce_nested_field(self.atom_encoder, AtomAttentionConfig)
-
-
-@dataclass
-class DiffusionStructureHeadConfig:
-    """Training and inference schedules for coordinate denoising."""
-
-    diffusion_module: DiffusionModuleConfig = field(default_factory=DiffusionModuleConfig)
-    distogram_bins: int = field(default=128)
-    train_noise_log_mean: float = field(default=-1.2)
-    train_noise_log_std: float = field(default=1.5)
-    gamma_0: float = field(default=0.605)
-    gamma_min: float = field(default=1.107)
-    noise_scale: float = field(default=0.0)
-    step_scale: float = field(default=1.0)
-    inference_s_max: float = field(default=160.0)
-    inference_s_min: float = field(default=4e-4)
-    inference_p: float = field(default=8.0)
-    inference_num_steps: int = field(default=68)
-
-    def __post_init__(self) -> None:
-        self.diffusion_module = _coerce_nested_field(self.diffusion_module, DiffusionModuleConfig)
-
-
-@dataclass
-class ConfidenceHeadConfig:
-    """Confidence-bin definitions and the compact confidence trunk."""
-
-    enabled: bool = field(default=True)
-    num_plddt_bins: int = field(default=50)
-    num_pde_bins: int = field(default=64)
-    num_pae_bins: int = field(default=64)
-    min_dist: float = field(default=2.0)
-    max_dist: float = field(default=52.0)
-    distogram_bins: int = field(default=128)
-    folding_trunk: FoldingTrunkConfig = field(
-        default_factory=lambda: FoldingTrunkConfig(n_layers=4)
-    )
-
-    def __post_init__(self) -> None:
-        self.folding_trunk = _coerce_nested_field(self.folding_trunk, FoldingTrunkConfig)
-
-
-@dataclass
-class MSAEncoderConfig:
-    """Optional multiple-sequence-alignment encoder settings."""
-
-    enabled: bool = field(default=False)
-    d_msa: int = field(default=128)
-    d_hidden: int = field(default=32)
-    n_layers: int = field(default=4)
-    n_heads_msa: int = field(default=8)
-    msa_head_width: int = field(default=32)
-
-
-@dataclass
-class LMEncoderConfig:
-    """Release-model pair encoder derived from language-model states."""
-
-    enabled: bool = field(default=True)
-    n_layers: int = field(default=4)
-    lm_dropout: float = field(default=0.25)
-    per_loop_lm_dropout: bool = field(default=True)
-
-
-@dataclass
-class ParcaeConfig:
-    """Release-model diffusion-loop scheduler settings."""
-
-    enabled: bool = field(default=True)
-    poisson_mean: float = field(default=3.0)
-    min_steps: int = field(default=1)
-    max_steps: int | None = field(default=6)
-    coda_n_layers: int = field(default=2)
-
-
-_SCALAR_DEFAULTS: tuple[tuple[str, Any], ...] = (
-    ("d_single", 384),
-    ("d_pair", 256),
-    ("n_relative_residx_bins", 32),
-    ("n_relative_chain_bins", 2),
-    ("num_loops", 10),
-    ("num_diffusion_samples", 8),
-    ("disable_msa_features", False),
-    ("lm_dropout", 0.0),
-    ("force_lm_dropout_during_inference", False),
-    ("lm_mask_pct", 0.0),
-    ("lm_d_model", 2560),
-    ("lm_num_layers", 80),
-)
-_NESTED_CONFIGS = (
-    ("inputs", InputsEmbedderConfig),
-    ("folding_trunk", FoldingTrunkConfig),
-    ("structure_head", DiffusionStructureHeadConfig),
-    ("confidence_head", ConfidenceHeadConfig),
-    ("msa_encoder", MSAEncoderConfig),
-    ("parcae", ParcaeConfig),
-    ("lm_encoder", LMEncoderConfig),
-)
-
-
-class ESMFold2Config(PretrainedConfig):
-    """Serializable ESMFold2 architecture, runtime, and precision settings."""
-
-    model_type = "esmfold2"
-    has_no_defaults_at_init = True
-
-    def __init__(self, **kwargs: Any) -> None:
-        legacy_backend = normalize_esmc_attention_implementation(kwargs.get("esmc_attn_backend"))
-        requested_backend = normalize_esmc_attention_implementation(
-            kwargs.get("attn_implementation")
-        )
-        resolved_backend = requested_backend or legacy_backend
-        kwargs["attn_implementation"] = resolved_backend
-        super().__init__(**kwargs)
-
-        self.type = kwargs.get("type", "release")
-        if self.type not in {"experimental", "release"}:
-            raise ValueError(
-                f"ESMFold2Config.type must be 'release' or 'experimental', got {self.type!r}"
-            )
-
-        for name, default in _SCALAR_DEFAULTS:
-            setattr(self, name, kwargs.get(name, default))
-
-        _official_esmc_repo, default_esmc_repo = _esmc_backbone_checkpoint_ids()
-        self.esmc_id = normalize_esmc_id(kwargs.get("esmc_id", default_esmc_repo))
-        self.esmc_attn_backend = resolved_backend
-        self.esmc_precision = str(kwargs.get("esmc_precision", "auto"))
-        if self.esmc_precision not in _ESMC_PRECISIONS:
-            raise ValueError(
-                "esmc_precision must be 'auto', 'bf16', 'fp32', or 'fp8', "
-                f"got {self.esmc_precision!r}."
-            )
-
-        for name, config_type in _NESTED_CONFIGS:
-            setattr(self, name, _nested_config(kwargs.get(name), config_type))
-        if not isinstance(self.msa_encoder.enabled, bool):
-            raise TypeError("msa_encoder.enabled must be a boolean.")
-        declared_msa_conditioning = kwargs.get("msa_conditioning")
-        if "msa_conditioning" in kwargs and not isinstance(declared_msa_conditioning, bool):
-            raise TypeError("msa_conditioning must be a boolean when provided.")
-        self.msa_conditioning = (
-            self.msa_encoder.enabled
-            if "msa_conditioning" not in kwargs
-            else declared_msa_conditioning
-        )
-        if self.msa_conditioning != self.msa_encoder.enabled:
-            raise ValueError(
-                "msa_conditioning must match msa_encoder.enabled; received "
-                f"{self.msa_conditioning!r} and {self.msa_encoder.enabled!r}."
-            )
-        self.msa_encoder_overwrite = bool(kwargs.get("msa_encoder_overwrite", True))
-
-        self.classifier_train_scope = str(kwargs.get("classifier_train_scope", "probe"))
-        if self.classifier_train_scope not in {"probe", "projection"}:
-            raise ValueError(
-                "classifier_train_scope must be 'probe' or 'projection', "
-                f"got {self.classifier_train_scope!r}."
-            )
-        self.classifier_probe_hidden_size = int(
-            kwargs.get("classifier_probe_hidden_size", 512)
-        )
-        self.classifier_probe_num_heads = int(kwargs.get("classifier_probe_num_heads", 4))
-        self.classifier_probe_dropout = float(kwargs.get("classifier_probe_dropout", 0.1))
-        self.classifier_hidden_size = int(kwargs.get("classifier_hidden_size", 4096))
-        self.classifier_dropout = float(kwargs.get("classifier_dropout", 0.2))
-        self.classifier_pooling_types = list(
-            kwargs.get("classifier_pooling_types", ["mean"])
-        )
-
-    def to_dict(self) -> dict[str, Any]:
-        output = cast(dict[str, Any], super().to_dict())
-        for name, _config_type in _NESTED_CONFIGS:
-            output[name] = asdict(getattr(self, name))
-        return output
-
-
-__all__ = [
-    "ESMFold2Config",
-    "LMEncoderConfig",
-    "MSAEncoderConfig",
-    "ParcaeConfig",
-    "normalize_esmc_attention_implementation",
-    "normalize_esmc_id",
-]
+# Copyright 2026 Biohub. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Configuration schema for release and experimental ESMFold2 checkpoints."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+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"})
+
+
+def _esmc_backbone_checkpoint_ids() -> tuple[str, str]:
+    """Return the manifest-pinned official and FastPLMs ESMC repositories."""
+
+    from fastplms.registry import RegistryError, get_model_registry
+
+    registry = get_model_registry()
+    family = registry.families["esmfold2"]
+    if family.backbone_model is None:
+        raise RegistryError("families.esmfold2 must declare backbone_model.")
+    backbone = registry[family.backbone_model]
+    return backbone.official.repo_id, backbone.fast.repo_id
+
+
+def normalize_esmc_id(esmc_id: str) -> str:
+    """Resolve an official ESMC identifier to its FastPLMs checkpoint mirror."""
+
+    from fastplms.registry import get_model_registry
+
+    official_repo, fast_repo = _esmc_backbone_checkpoint_ids()
+    if esmc_id == official_repo:
+        return fast_repo
+    registry = get_model_registry()
+    for spec in registry.by_family("esmfold2"):
+        if spec.backbone is not None and spec.backbone_model is not None:
+            backbone = registry[spec.backbone_model]
+            if esmc_id in {spec.backbone.repo_id, backbone.official.repo_id}:
+                return backbone.fast.repo_id
+    return esmc_id
+
+
+def normalize_esmc_attention_implementation(
+    implementation: str | dict[str, str] | None,
+) -> str | None:
+    """Validate the ESMC backend and translate the historical ``flex`` name."""
+
+    if isinstance(implementation, dict):
+        if tuple(implementation) != ("",):
+            raise ValueError(
+                "ESMFold2 has one ESMC attention backbone; use a string or {'': implementation}."
+            )
+        implementation = 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(
+            f"Unsupported ESMFold2 attention implementation {canonical!r}; "
+            f"expected one of {expected}."
+        )
+    return canonical
+
+
+NestedConfig = TypeVar("NestedConfig")
+
+
+def _nested_config(value: Any, config_type: type[NestedConfig]) -> NestedConfig:
+    if isinstance(value, config_type):
+        return value
+    return config_type(**value) if isinstance(value, dict) else config_type()
+
+
+def _coerce_nested_field(
+    value: NestedConfig | dict[str, Any], config_type: type[NestedConfig]
+) -> NestedConfig:
+    """Convert serialized nested dictionaries while retaining supplied objects."""
+
+    return config_type(**value) if isinstance(value, dict) else value
+
+
+@dataclass
+class AtomAttentionConfig:
+    """Sliding-window atom attention and three-dimensional RoPE settings."""
+
+    d_atom: int = field(default=128)
+    d_token: int = field(default=768)
+    n_blocks: int = field(default=3)
+    n_heads: int = field(default=4)
+    swa_window_size: int = field(default=128)
+    expansion_ratio: int = field(default=2)
+    spatial_rope_base_frequency: float = field(default=20.0)
+    n_spatial_rope_pairs_per_axis: int = field(default=2)
+    n_uid_rope_pairs: int = field(default=10)
+    uid_rope_base_frequency: float = field(default=10000.0)
+
+
+@dataclass
+class DiffusionModuleConfig:
+    """Dimensions and depth of the coordinate diffusion network."""
+
+    sigma_data: float = field(default=16.0)
+    c_atom: int = field(default=128)
+    c_token: int = field(default=768)
+    c_z: int = field(default=256)
+    c_s_inputs: int = field(default=451)
+    fourier_dim: int = field(default=256)
+    relpos_r_max: int = field(default=32)
+    relpos_s_max: int = field(default=2)
+    atom_num_blocks: int = field(default=3)
+    atom_num_heads: int = field(default=4)
+    token_num_blocks: int = field(default=12)
+    token_num_heads: int = field(default=16)
+    transition_multiplier: int = field(default=2)
+
+
+@dataclass
+class FoldingTrunkConfig:
+    """Iterative pair/single trunk dimensions."""
+
+    n_layers: int = field(default=24)
+    n_heads: int = field(default=8)
+    dropout: float = field(default=0.0)
+
+
+@dataclass
+class InputsEmbedderConfig:
+    """Input feature width and atom encoder settings."""
+
+    d_inputs: int = field(default=451)
+    atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig)
+
+    def __post_init__(self) -> None:
+        self.atom_encoder = _coerce_nested_field(self.atom_encoder, AtomAttentionConfig)
+
+
+@dataclass
+class DiffusionStructureHeadConfig:
+    """Training and inference schedules for coordinate denoising."""
+
+    diffusion_module: DiffusionModuleConfig = field(default_factory=DiffusionModuleConfig)
+    distogram_bins: int = field(default=128)
+    train_noise_log_mean: float = field(default=-1.2)
+    train_noise_log_std: float = field(default=1.5)
+    gamma_0: float = field(default=0.605)
+    gamma_min: float = field(default=1.107)
+    noise_scale: float = field(default=0.0)
+    step_scale: float = field(default=1.0)
+    inference_s_max: float = field(default=160.0)
+    inference_s_min: float = field(default=4e-4)
+    inference_p: float = field(default=8.0)
+    inference_num_steps: int = field(default=68)
+
+    def __post_init__(self) -> None:
+        self.diffusion_module = _coerce_nested_field(self.diffusion_module, DiffusionModuleConfig)
+
+
+@dataclass
+class ConfidenceHeadConfig:
+    """Confidence-bin definitions and the compact confidence trunk."""
+
+    enabled: bool = field(default=True)
+    num_plddt_bins: int = field(default=50)
+    num_pde_bins: int = field(default=64)
+    num_pae_bins: int = field(default=64)
+    min_dist: float = field(default=2.0)
+    max_dist: float = field(default=52.0)
+    distogram_bins: int = field(default=128)
+    folding_trunk: FoldingTrunkConfig = field(
+        default_factory=lambda: FoldingTrunkConfig(n_layers=4)
+    )
+
+    def __post_init__(self) -> None:
+        self.folding_trunk = _coerce_nested_field(self.folding_trunk, FoldingTrunkConfig)
+
+
+@dataclass
+class MSAEncoderConfig:
+    """Optional multiple-sequence-alignment encoder settings."""
+
+    enabled: bool = field(default=False)
+    d_msa: int = field(default=128)
+    d_hidden: int = field(default=32)
+    n_layers: int = field(default=4)
+    n_heads_msa: int = field(default=8)
+    msa_head_width: int = field(default=32)
+
+
+@dataclass
+class LMEncoderConfig:
+    """Release-model pair encoder derived from language-model states."""
+
+    enabled: bool = field(default=True)
+    n_layers: int = field(default=4)
+    lm_dropout: float = field(default=0.25)
+    per_loop_lm_dropout: bool = field(default=True)
+
+
+@dataclass
+class ParcaeConfig:
+    """Release-model diffusion-loop scheduler settings."""
+
+    enabled: bool = field(default=True)
+    poisson_mean: float = field(default=3.0)
+    min_steps: int = field(default=1)
+    max_steps: int | None = field(default=6)
+    coda_n_layers: int = field(default=2)
+
+
+_SCALAR_DEFAULTS: tuple[tuple[str, Any], ...] = (
+    ("d_single", 384),
+    ("d_pair", 256),
+    ("n_relative_residx_bins", 32),
+    ("n_relative_chain_bins", 2),
+    ("num_loops", 10),
+    ("num_diffusion_samples", 8),
+    ("disable_msa_features", False),
+    ("lm_dropout", 0.0),
+    ("force_lm_dropout_during_inference", False),
+    ("lm_mask_pct", 0.0),
+    ("lm_d_model", 2560),
+    ("lm_num_layers", 80),
+)
+_NESTED_CONFIGS = (
+    ("inputs", InputsEmbedderConfig),
+    ("folding_trunk", FoldingTrunkConfig),
+    ("structure_head", DiffusionStructureHeadConfig),
+    ("confidence_head", ConfidenceHeadConfig),
+    ("msa_encoder", MSAEncoderConfig),
+    ("parcae", ParcaeConfig),
+    ("lm_encoder", LMEncoderConfig),
+)
+
+
+class ESMFold2Config(PretrainedConfig):
+    """Serializable ESMFold2 architecture, runtime, and precision settings."""
+
+    model_type = "esmfold2"
+    has_no_defaults_at_init = True
+
+    def __init__(self, **kwargs: Any) -> None:
+        legacy_backend = normalize_esmc_attention_implementation(kwargs.get("esmc_attn_backend"))
+        requested_backend = normalize_esmc_attention_implementation(
+            kwargs.get("attn_implementation")
+        )
+        resolved_backend = requested_backend or legacy_backend
+        kwargs["attn_implementation"] = resolved_backend
+        super().__init__(**kwargs)
+
+        self.type = kwargs.get("type", "release")
+        if self.type not in {"experimental", "release"}:
+            raise ValueError(
+                f"ESMFold2Config.type must be 'release' or 'experimental', got {self.type!r}"
+            )
+
+        for name, default in _SCALAR_DEFAULTS:
+            setattr(self, name, kwargs.get(name, default))
+
+        _official_esmc_repo, default_esmc_repo = _esmc_backbone_checkpoint_ids()
+        self.esmc_id = normalize_esmc_id(kwargs.get("esmc_id", default_esmc_repo))
+        self.esmc_attn_backend = resolved_backend
+        self.esmc_precision = str(kwargs.get("esmc_precision", "auto"))
+        if self.esmc_precision not in _ESMC_PRECISIONS:
+            raise ValueError(
+                "esmc_precision must be 'auto', 'bf16', 'fp32', or 'fp8', "
+                f"got {self.esmc_precision!r}."
+            )
+
+        for name, config_type in _NESTED_CONFIGS:
+            setattr(self, name, _nested_config(kwargs.get(name), config_type))
+        if not isinstance(self.msa_encoder.enabled, bool):
+            raise TypeError("msa_encoder.enabled must be a boolean.")
+        declared_msa_conditioning = kwargs.get("msa_conditioning")
+        if "msa_conditioning" in kwargs and not isinstance(declared_msa_conditioning, bool):
+            raise TypeError("msa_conditioning must be a boolean when provided.")
+        self.msa_conditioning = (
+            self.msa_encoder.enabled
+            if "msa_conditioning" not in kwargs
+            else declared_msa_conditioning
+        )
+        if self.msa_conditioning != self.msa_encoder.enabled:
+            raise ValueError(
+                "msa_conditioning must match msa_encoder.enabled; received "
+                f"{self.msa_conditioning!r} and {self.msa_encoder.enabled!r}."
+            )
+        self.msa_encoder_overwrite = bool(kwargs.get("msa_encoder_overwrite", True))
+
+        self.classifier_train_scope = str(kwargs.get("classifier_train_scope", "probe"))
+        if self.classifier_train_scope not in {"probe", "projection"}:
+            raise ValueError(
+                "classifier_train_scope must be 'probe' or 'projection', "
+                f"got {self.classifier_train_scope!r}."
+            )
+        self.classifier_probe_hidden_size = int(
+            kwargs.get("classifier_probe_hidden_size", 512)
+        )
+        self.classifier_probe_num_heads = int(kwargs.get("classifier_probe_num_heads", 4))
+        self.classifier_probe_dropout = float(kwargs.get("classifier_probe_dropout", 0.1))
+        self.classifier_hidden_size = int(kwargs.get("classifier_hidden_size", 4096))
+        self.classifier_dropout = float(kwargs.get("classifier_dropout", 0.2))
+        self.classifier_pooling_types = list(
+            kwargs.get("classifier_pooling_types", ["mean"])
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        output = cast(dict[str, Any], super().to_dict())
+        for name, _config_type in _NESTED_CONFIGS:
+            output[name] = asdict(getattr(self, name))
+        return output
+
+
+__all__ = [
+    "ESMFold2Config",
+    "LMEncoderConfig",
+    "MSAEncoderConfig",
+    "ParcaeConfig",
+    "normalize_esmc_attention_implementation",
+    "normalize_esmc_id",
+]

fastplms/models/esmfold2/esmfold2_msa.py

@@ -239,15 +239,18 @@ class MSA(SequentialDataclass):
             raise ValueError("MSA requires at least one aligned sequence.")
         if any(not isinstance(entry, FastaEntry) for entry in self.entries):
             raise TypeError("Every MSA entry must be a FastaEntry.")
-        expected_length = len(self.entries[0].sequence)
+        # Rows align on match columns. A raw A3M row also carries lowercase and
+        # "." insertions, which the pinned Biohub MSA accepts and paired-MSA
+        # construction decodes into deletion counts, so they are not columns.
+        expected_length = len(remove_insertions_from_sequence(self.entries[0].sequence))
         if expected_length == 0:
             raise ValueError("MSA sequences must be non-empty.")
         for row, entry in enumerate(self.entries[1:], start=1):
-            if len(entry.sequence) != expected_length:
+            row_length = len(remove_insertions_from_sequence(entry.sequence))
+            if row_length != expected_length:
                 raise ValueError(
                     "MSA row length mismatch: "
-                    f"row 0 has {expected_length} columns, row {row} has "
-                    f"{len(entry.sequence)}."
+                    f"row 0 has {expected_length} columns, row {row} has {row_length}."
                 )
         deletions = self.deletions
         if deletions is not None and not isinstance(deletions, np.ndarray):

fastplms/models/esmfold2/modeling_esmfold2_common.py

@@ -12,7 +12,7 @@ from __future__ import annotations
 import importlib
 from functools import partial
 from importlib.util import find_spec
-from typing import ClassVar, cast
+from typing import Any, ClassVar, cast
 
 import torch
 import torch.nn as nn
@@ -39,6 +39,13 @@ except (AttributeError, ImportError):
     _cue_tri_mul = None  # type: ignore[assignment]
     CUE_AVAILABLE = False
 
+# PyTorch ships this variable-length FlashAttention entry point, so windowed
+# atom attention needs no extra package and no compilation.
+try:
+    from torch.nn.attention.varlen import varlen_attn as _varlen_attn
+except ImportError:
+    _varlen_attn = None  # type: ignore[assignment]
+
 # Biohub ships optional source-built Triton helpers. FastPLMs does not bundle or
 # compile them; these placeholders retain checkpoint-compatible control flow
 # while the portable PyTorch path remains flat and self-contained.
@@ -51,6 +58,9 @@ TRITON_KERNELS_AVAILABLE = False
 BACKEND_FUSED = "fused"
 BACKEND_CUEQ = "cuequivariance"
 _VALID_BACKENDS = (None, BACKEND_FUSED, BACKEND_CUEQ)
+ATOM_ATTENTION_DENSE = "dense"
+ATOM_ATTENTION_WINDOWED = "windowed"
+_VALID_ATOM_ATTENTION = (ATOM_ATTENTION_DENSE, ATOM_ATTENTION_WINDOWED)
 MSA_CONDITIONING_INPUT_NAMES = (
     "msa",
     "msa_attention_mask",
@@ -86,6 +96,18 @@ def validate_kernel_backend(backend: str | None) -> None:
         )
 
 
+def validate_atom_attention(mode: str) -> None:
+    """Fail before mutating modules when a named atom attention mode cannot execute."""
+
+    if mode not in _VALID_ATOM_ATTENTION:
+        raise ValueError(f"atom attention must be one of {_VALID_ATOM_ATTENTION}, got {mode!r}")
+    if mode == ATOM_ATTENTION_WINDOWED and _varlen_attn is None:
+        raise RuntimeError(
+            "atom attention 'windowed' requires torch.nn.attention.varlen, which this "
+            "PyTorch build does not provide."
+        )
+
+
 def validate_msa_conditioning_inputs(
     config: ESMFold2Config,
     *,
@@ -827,18 +849,29 @@ class TriangleMultiplicativeBlock(nn.Module):
     def _triangular_contract_chunked(
         self, left_stream: Tensor, right_stream: Tensor, chunk_size: int
     ) -> Tensor:
-        """Compute the triangular einsum in chunks along the output i-dimension."""
-        length = left_stream.shape[1] if self.flow == "outgoing" else left_stream.shape[2]
+        """Compute the triangular einsum in chunks along the output i-dimension.
+
+        ``torch.einsum`` lowers each chunk to one batched matrix product over
+        ``(b, d)``, and to do so it copies the whole right stream into that
+        layout once per chunk. That copy is made once here. Each chunk then runs
+        the same product on the same values, so the result is bitwise identical
+        to a per-chunk einsum. ``forward`` has already cast the streams to the
+        autocast dtype, which the product would otherwise do once per chunk.
+        """
+        if self.flow == "outgoing":
+            left_rows = left_stream.permute(0, 3, 1, 2)  # (b, i, k, d) -> (b, d, i, k)
+            right_columns = right_stream.permute(0, 3, 2, 1)  # (b, j, k, d) -> (b, d, k, j)
+        else:
+            left_rows = left_stream.permute(0, 3, 2, 1)  # (b, k, i, d) -> (b, d, i, k)
+            right_columns = right_stream.permute(0, 3, 1, 2)  # (b, k, j, d) -> (b, d, k, j)
+        batch_size, channels, length, inner = left_rows.shape
+        right_columns = right_columns.reshape(batch_size * channels, inner, -1)  # (b * d, k, j)
         chunks = []
         for start in range(0, length, chunk_size):
-            end = min(start + chunk_size, length)
-            if self.flow == "outgoing":
-                chunk = torch.einsum(self._einsum_equation, left_stream[:, start:end], right_stream)
-            else:
-                chunk = torch.einsum(
-                    self._einsum_equation, left_stream[:, :, start:end], right_stream
-                )
-            chunks.append(chunk)
+            rows = left_rows[:, :, start : start + chunk_size]  # (b, d, i_c, k)
+            product = torch.bmm(rows.reshape(batch_size * channels, -1, inner), right_columns)
+            product = product.view(batch_size, channels, rows.shape[2], -1)  # (b, d, i_c, j)
+            chunks.append(product.permute(0, 2, 3, 1))  # (b, i_c, j, d)
         return torch.cat(chunks, dim=1)
 
     def forward(self, pair_grid: Tensor, visibility: Tensor | None = None) -> Tensor:
@@ -862,20 +895,34 @@ class TriangleMultiplicativeBlock(nn.Module):
                 eps=_EPS,
             )
 
+        # Every tensor below is as large as the pair representation or larger, and this
+        # block sets the peak memory of a fold. Each name is dropped once it is dead, so
+        # the allocator can reuse its buffer. No value changes.
         normalized_grid = self.norm_start(pair_grid)
-        bundled = self.proj_bundle(normalized_grid)
+        bundled = self.proj_bundle(normalized_grid)  # (b, l, l, 4 * d)
         signal, gate_logits = bundled.split(2 * self.latent_channels, dim=-1)
-        routed = signal * torch.sigmoid(gate_logits)
+        routed = signal * torch.sigmoid(gate_logits)  # (b, l, l, 2 * d)
+        # The two views would keep the whole projection alive.
+        del bundled, signal, gate_logits
         routed = routed * visibility.unsqueeze(-1)
 
-        left_stream, right_stream = routed.float().chunk(2, dim=-1)
+        left_stream, right_stream = routed.float().chunk(2, dim=-1)  # each (b, l, l, d)
+        if torch.is_autocast_enabled(left_stream.device.type):
+            # The contraction is an autocast operation. Casting its inputs here, as it
+            # would, lets the full-precision product go before the contraction runs.
+            autocast_dtype = torch.get_autocast_dtype(left_stream.device.type)
+            left_stream = left_stream.to(autocast_dtype)
+            right_stream = right_stream.to(autocast_dtype)
+        del routed
         if self._chunk_size is not None:
             contracted = self._triangular_contract_chunked(
                 left_stream, right_stream, self._chunk_size
             )
         else:
             contracted = self._triangular_contract(left_stream, right_stream)
+        del left_stream, right_stream
         mixed = self.proj_emit(self.norm_mix(contracted))
+        del contracted
         output_gate = torch.sigmoid(self.proj_gate(normalized_grid))
         return mixed * output_gate
 
@@ -1315,7 +1362,11 @@ class SwiGLU(nn.Module):
     def forward(self, x: Tensor) -> Tensor:
         x12 = self.w12(x)
         x1, x2 = x12.split(self.hidden_features, dim=-1)
-        hidden = F.silu(x1) * x2
+        hidden = F.silu(x1)
+        # Without autograd the product can reuse the activation's buffer. On a pair tensor
+        # that buffer is twice the pair representation. The values are the same either way.
+        hidden = hidden * x2 if torch.is_grad_enabled() else hidden.mul_(x2)
+        del x12, x1, x2
         return self.w3(hidden)
 
 
@@ -1447,11 +1498,58 @@ class SWA3DRoPEAttention(nn.Module):
         self.head_dim = d_model // n_heads
         self.scale = self.head_dim**-0.5
         self.half_window = half_window
+        self._atom_attention = ATOM_ATTENTION_DENSE
 
         self.Wqkv = nn.Linear(d_model, 3 * d_model, bias=False)
         self.out_proj = nn.Linear(d_model, d_model, bias=False)
         self.gate_proj = nn.Linear(d_model, d_model, bias=False)
 
+    def set_atom_attention(self, mode: str) -> None:
+        validate_atom_attention(mode)
+        self._atom_attention = mode
+
+    def _windowed_attention(
+        self, q: Tensor, k: Tensor, v: Tensor, attention_params: tuple[Any, ...]
+    ) -> Tensor:
+        """Attend to at most ``half_window`` real atoms on each side.
+
+        This is the attention the official model runs when flash-attn is
+        installed. Window offsets count unpadded atoms within one sample, so a
+        padded slot is never a key, and its output is zero.
+        """
+        # q, k, v: (b, n_atoms, h, d_h)
+        if not q.is_cuda:
+            raise RuntimeError("atom attention 'windowed' requires CUDA tensors.")
+        batch_size, n_atoms = q.shape[:2]
+        # indices: (t,) flat positions of real atoms; cu_seqlens: (b + 1,) int32 row offsets.
+        indices, cu_seqlens, max_seqlen = attention_params[2:5]
+        flat_shape = (batch_size * n_atoms, self.n_heads, self.head_dim)
+        q, k, v = q.reshape(flat_shape), k.reshape(flat_shape), v.reshape(flat_shape)
+        has_padding = indices.shape[0] != batch_size * n_atoms
+        if has_padding:
+            q, k, v = q[indices], k[indices], v[indices]  # each (t, h, d_h)
+        # Without an auxiliary request the call returns the attended values alone.
+        attended = cast(
+            Tensor,
+            _varlen_attn(
+                q,
+                k,
+                v,
+                cu_seqlens,
+                cu_seqlens,
+                max_seqlen,
+                max_seqlen,
+                scale=self.scale,
+                window_size=(self.half_window, self.half_window),
+            ),
+        )  # (t, h, d_h)
+        if has_padding:
+            out = attended.new_zeros(flat_shape)  # (b * n_atoms, h, d_h)
+            out[indices] = attended
+        else:
+            out = attended
+        return out.view(batch_size, n_atoms, self.n_heads, self.head_dim)
+
     def forward(self, x: Tensor, attention_params: tuple) -> Tensor:
         batch_size, n_atoms = x.shape[:2]
         cos, sin = attention_params[0], attention_params[1]
@@ -1472,12 +1570,15 @@ class SWA3DRoPEAttention(nn.Module):
         # ESMFold2 does not advertise FlashAttention. Keep this atom path on
         # PyTorch. Models that advertise FlashAttention dispatch through the
         # precompiled Hugging Face kernels interface in fastplms.attention.
-        q_t = q.transpose(1, 2)
-        k_t = k.transpose(1, 2)
-        v_t = v.transpose(1, 2)
-        attn = torch.matmul(q_t, k_t.transpose(-2, -1)) * self.scale
-        attn = F.softmax(attn, dim=-1)
-        out = torch.matmul(attn, v_t).transpose(1, 2)
+        if self._atom_attention == ATOM_ATTENTION_WINDOWED:
+            out = self._windowed_attention(q, k, v, attention_params)
+        else:
+            q_t = q.transpose(1, 2)
+            k_t = k.transpose(1, 2)
+            v_t = v.transpose(1, 2)
+            attn = torch.matmul(q_t, k_t.transpose(-2, -1)) * self.scale
+            attn = F.softmax(attn, dim=-1)
+            out = torch.matmul(attn, v_t).transpose(1, 2)
 
         out = out.to(input_dtype).reshape(  # type: ignore[union-attr]
             batch_size, n_atoms, -1
@@ -1934,7 +2035,14 @@ class AttentionPairBias(nn.Module):
         beta: Tensor | float = 0.0,
         attention_mask: Tensor | None = None,
         num_diffusion_samples: int = 1,
+        step_cache: dict[str, Tensor] | None = None,
     ) -> Tensor:
+        """Attend over tokens with a per-head bias projected from the pair tensor.
+
+        ``step_cache`` belongs to this block for one ``sample`` call. The sampler
+        conditions every denoising step on the same ``z``, so the PyTorch path
+        projects the bias on the first step and reuses that tensor afterwards.
+        """
         bsz, n_queries, d_model = a.shape
 
         x = self.adaln(a, s) if s is not None else self.pre_norm(a)
@@ -1946,9 +2054,20 @@ class AttentionPairBias(nn.Module):
         k = k.view(bsz, n_keys, self.num_heads, self.head_dim)
         v = v.view(bsz, n_keys, self.num_heads, self.head_dim)
 
-        # Expand z for num_diffusion_samples
-        if z.dim() == 4 and z.shape[0] != bsz and num_diffusion_samples > 1:
-            z = z.repeat_interleave(num_diffusion_samples, dim=0)
+        use_fused_kernel = self._can_use_fused_pair_bias(z, n_queries, beta)
+        use_cueq_kernel = not use_fused_kernel and self._can_use_cueq_pair_bias(z, n_queries, beta)
+        cached_pair_bias = None
+        if step_cache is not None and not use_fused_kernel and not use_cueq_kernel:
+            cached_pair_bias = step_cache.get("pair_bias")
+
+        # Expand z for num_diffusion_samples, unless its projection is already cached.
+        if (
+            cached_pair_bias is None
+            and z.dim() == 4
+            and z.shape[0] != bsz
+            and num_diffusion_samples > 1
+        ):
+            z = z.repeat_interleave(num_diffusion_samples, dim=0)  # (b * samples, n, n, d_pair)
         if (
             attention_mask is not None
             and attention_mask.shape[0] != bsz
@@ -1956,7 +2075,7 @@ class AttentionPairBias(nn.Module):
         ):
             attention_mask = attention_mask.repeat_interleave(num_diffusion_samples, dim=0)
 
-        if self._can_use_fused_pair_bias(z, n_queries, beta):
+        if use_fused_kernel:
             kernel_mask = (
                 attention_mask
                 if attention_mask is not None
@@ -1990,7 +2109,7 @@ class AttentionPairBias(nn.Module):
                 out = torch.sigmoid(self.out_gate(s)) * out
             return out
 
-        if self._can_use_cueq_pair_bias(z, n_queries, beta):
+        if use_cueq_kernel:
             kernel_mask = (
                 attention_mask
                 if attention_mask is not None
@@ -2018,7 +2137,14 @@ class AttentionPairBias(nn.Module):
 
             logits = torch.einsum("... i h d, ... j h d -> ... i j h", q, k) * self.scale
 
-            pair_bias = self.pair_bias_proj(self.pair_norm(z)) if z.dim() == 4 else z.unsqueeze(-1)
+            if cached_pair_bias is not None:
+                pair_bias = cached_pair_bias  # (b * samples, n, n, h)
+            elif z.dim() == 4:
+                pair_bias = self.pair_bias_proj(self.pair_norm(z))  # (b * samples, n, n, h)
+                if step_cache is not None:
+                    step_cache["pair_bias"] = pair_bias
+            else:
+                pair_bias = z.unsqueeze(-1)  # (b * samples, n, n, 1), a precomputed bias
             logits = logits + pair_bias.to(dtype=logits.dtype)
 
             if attention_mask is not None:
@@ -2136,10 +2262,22 @@ class DiffusionTransformer(nn.Module):
         attention_mask: Tensor | None = None,
         num_diffusion_samples: int = 1,
         return_intermediates: bool = False,
+        inference_cache: dict[str, Any] | None = None,
     ) -> tuple[Tensor, list[Tensor]]:
+        """Run the token blocks.
+
+        ``inference_cache`` must span only calls that share ``z``, as one
+        ``sample`` call does; each block then keeps its pair bias across steps.
+        """
         intermediates: list[Tensor] = []
+        block_caches: dict[int, dict[str, Tensor]] | None = None
+        if inference_cache is not None:
+            block_caches = inference_cache.setdefault("token_pair_bias", {})
         x = a
-        for attn, transition in zip(self.attn_blocks, self.transition_blocks, strict=True):
+        for block_index, (attn, transition) in enumerate(
+            zip(self.attn_blocks, self.transition_blocks, strict=True)
+        ):
+            step_cache = None if block_caches is None else block_caches.setdefault(block_index, {})
             x = x + attn(
                 x,
                 s,
@@ -2147,6 +2285,7 @@ class DiffusionTransformer(nn.Module):
                 beta,
                 attention_mask=attention_mask,
                 num_diffusion_samples=num_diffusion_samples,
+                step_cache=step_cache,
             )
             x = x + transition(x, s)
             if return_intermediates:
@@ -2409,6 +2548,7 @@ class DiffusionModule(nn.Module):
             beta=0.0,
             attention_mask=token_attention_mask,
             num_diffusion_samples=num_diffusion_samples,
+            inference_cache=inference_cache,
         )
 
         # Step 6: token norm

fastplms/models/esmfold2/modeling_esmfold2_experimental.py

@@ -51,6 +51,7 @@ from .modeling_esmfold2_common import (
     PairUpdateBlock,
     ResIdxAsymIdSymIdEntityIdEncoding,
     RowAttentionPooling,
+    SWA3DRoPEAttention,
     SwiGLUMLP,
     TriangleMultiplicativeUpdate,
     _categorical_mean,
@@ -59,6 +60,7 @@ from .modeling_esmfold2_common import (
     compute_lm_hidden_states,
     gather_rep_atom_coords,
     gather_token_to_atom,
+    validate_atom_attention,
     validate_kernel_backend,
     validate_msa_conditioning_inputs,
     validate_prepared_auxiliary_inputs,
@@ -511,6 +513,22 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         self.structure_head.set_kernel_backend(backend)
         self._kernel_backend = backend
 
+    def set_atom_attention(self, mode: str) -> None:
+        """Select how atoms attend to each other in the atom encoders and decoder.
+
+        Args:
+            mode: ``"dense"`` (default) lets every atom attend to every atom slot,
+                as the official model does without flash-attn. ``"windowed"``
+                restricts each atom to 64 real neighbors on each side through
+                PyTorch's variable-length FlashAttention, as the official model
+                does with flash-attn. It needs CUDA, costs linear rather than
+                quadratic memory in the atom count, and changes numerical output.
+        """
+        validate_atom_attention(mode)
+        for module in self.modules():
+            if isinstance(module, SWA3DRoPEAttention):
+                module.set_atom_attention(mode)
+
     def set_chunk_size(self, chunk_size: int | None) -> None:
         self.folding_trunk.set_chunk_size(chunk_size)
         if self.confidence_head is not None:
@@ -606,7 +624,6 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
     def apply_torch_compile(self, mode: str = "fixed_seqlen", dynamic: bool | None = None) -> None:
         if dynamic is None:
             dynamic = mode == "dynamic_seqlen"
-        compile_kwargs: dict[str, bool] = {"dynamic": dynamic}
         compile_targets = (
             PairUpdateBlock,
             DiffusionTransformer,
@@ -616,7 +633,9 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
 
         def _maybe_compile(module: nn.Module) -> None:
             if isinstance(module, compile_targets):
-                module.forward = torch.compile(module.forward, **compile_kwargs)
+                module.forward = torch.compile(  # type: ignore[assignment]
+                    module.forward, dynamic=dynamic
+                )
 
         self.apply(_maybe_compile)
 

modeling_fastplms.py

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