Update FastPLMs files from Synthyra/FastPLMs 32d9951 (PR 52: ESMFold2 noise cap, DPLM2 parity)

README.md

@@ -79,7 +79,7 @@ python -m pip install -r \
 The FastPLMs implementation itself is embedded in the model repository.
 Transformers loads it through `trust_remote_code=True`.
 
-This model requires Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13.
+This model requires Python 3.12-3.14, PyTorch 2.14, and Transformers 5.17.
 
 The artifact requirements include the structure dependencies.
 
@@ -144,7 +144,7 @@ print(token_output.logits.shape)     # (b, l, 3)
 Install the training dependencies. Then attach LoRA to the loaded checkpoint:
 
 ```bash
-python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20"
+python -m pip install "datasets>=5.0" "peft>=0.21"
 ```
 
 ```python

fastplms/atomic_files.py

@@ -0,0 +1,65 @@
+"""Replace a file in one step: readers see the old bytes or the new bytes, never a partial write."""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import IO, Any
+
+
+def write_bytes_atomically(path: Path, payload: bytes, *, create_parent: bool = False) -> None:
+    """Write ``payload`` to a temporary file beside ``path``, flush it to disk, then rename it over ``path``.
+
+    A failed write removes the temporary file and leaves ``path`` untouched. The parent directory must
+    exist unless ``create_parent`` is true. The temporary file sits in the same directory so the rename
+    never crosses a file system.
+    """
+
+    with _staged_handle(path, "wb", create_parent=create_parent) as handle:
+        handle.write(payload)
+
+
+def write_text_atomically(
+    path: Path,
+    text: str,
+    *,
+    encoding: str = "utf-8",
+    newline: str | None = None,
+    create_parent: bool = False,
+) -> None:
+    """Write ``text`` as ``write_bytes_atomically`` does.
+
+    ``newline`` is the argument of ``open``: ``None`` translates ``"\\n"`` to the platform separator and
+    ``"\\n"`` writes it unchanged.
+    """
+
+    with _staged_handle(
+        path, "w", create_parent=create_parent, encoding=encoding, newline=newline
+    ) as handle:
+        handle.write(text)
+
+
+@contextmanager
+def _staged_handle(
+    path: Path, mode: str, *, create_parent: bool, **open_arguments: Any
+) -> Iterator[IO[Any]]:
+    """Yield a handle to a temporary file; on a clean exit flush it and rename it over ``path``."""
+
+    if create_parent:
+        path.parent.mkdir(parents=True, exist_ok=True)
+    descriptor, temporary_name = tempfile.mkstemp(
+        prefix=f".{path.name}.", suffix=".tmp", dir=path.parent
+    )
+    try:
+        with os.fdopen(descriptor, mode, **open_arguments) as handle:
+            yield handle
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temporary_name, path)
+    except BaseException:
+        Path(temporary_name).unlink(missing_ok=True)
+        raise

fastplms/attention/__init__.py

@@ -22,6 +22,7 @@ from ._core import (
     canonical_checkpoint_attention_backend,
     clear_flex_attention_caches,
     create_block_mask,
+    flash_kernel_unsupported_reason,
     flex_attention,
     get_attention_mask,
     get_attn_implementation,
@@ -66,6 +67,7 @@ __all__ = [
     "canonical_checkpoint_attention_backend",
     "clear_flex_attention_caches",
     "create_block_mask",
+    "flash_kernel_unsupported_reason",
     "flex_attention",
     "get_attention_mask",
     "get_attn_implementation",

fastplms/attention/_core.py

@@ -16,10 +16,15 @@ from dataclasses import dataclass
 from enum import Enum
 from threading import RLock
 from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, TypeVar
 from einops import rearrange
 from torch.nn import functional as F
 
-from ._kernel_lock import load_locked_kernel
+from ._kernel_lock import load_locked_kernel, locked_variant_for_this_system
+
+
+if TYPE_CHECKING:
+    from transformers import PretrainedConfig
 
 
 try:
@@ -30,12 +35,16 @@ except ImportError:
     BlockMask = None
 
 _MAX_FLEX_CACHE_ENTRIES = 128
-_compiled_flex_attention: OrderedDict[tuple, object] = OrderedDict()
-_flex_block_masks: OrderedDict[tuple, BlockMask] = OrderedDict()
+_compiled_flex_attention: OrderedDict[tuple[Any, ...], object] = OrderedDict()
+_flex_block_masks: OrderedDict[tuple[Any, ...], BlockMask] = OrderedDict()
 _flex_cache_lock = RLock()
 
+_CacheValue = TypeVar("_CacheValue")
+
 
-def _remember(cache: OrderedDict, key: tuple, value):
+def _remember(
+    cache: OrderedDict[tuple[Any, ...], _CacheValue], key: tuple[Any, ...], value: _CacheValue
+) -> _CacheValue:
     """Insert an item into a bounded least-recently-used cache."""
     cache[key] = value
     cache.move_to_end(key)
@@ -64,7 +73,7 @@ def _get_flex_attention_fn(
     shape: tuple[int, ...] | None = None,
     sequence_lengths: tuple[int, ...] | None = None,
     mask_semantics: str = "padding",
-):
+) -> Callable[..., Any] | None:
     """Return a compiled Flex callable for an explicit execution signature.
 
     Compilation depends on execution shape, device, dtype, and mask semantics.
@@ -116,15 +125,17 @@ def _get_flex_block_mask(
     because compiled Flex plans can specialize on it even though the pattern
     tensor itself is boolean or integer.
     """
+    # mask_pattern: (b, l), the padding mask of the call
+    # mask_mod: ((), (), (), ()) -> (), taking the 0-d batch, head, query and key-value indices
     if create_block_mask is None:
         raise RuntimeError(
             "'flex_attention' was requested, but torch.create_block_mask is unavailable."
         )
-    pattern = mask_pattern.detach().to(device=device).contiguous()  # mask_pattern.shape
+    pattern = mask_pattern.detach().to(device=device).contiguous()  # (b, l)
     # One device-to-host transfer is required for an exact cache identity. Use
     # the contiguous buffer directly instead of materializing one Python int
     # per byte, which is prohibitively expensive for long batched sequences.
-    host_pattern = pattern.to(device="cpu").contiguous()  # mask_pattern.shape
+    host_pattern = pattern.to(device="cpu").contiguous()  # (b, l)
     pattern_bytes = host_pattern.view(torch.uint8).numpy().tobytes(order="C")  # bytes
     cache_key = (
         str(device),
@@ -153,7 +164,7 @@ def _get_flex_block_mask(
 
 # Hugging Face `kernels` exposes slightly different APIs for FlashAttention 2
 # and 3. Detect the loaded variant once so every caller uses the same dispatch.
-def _infer_kernels_flash_variant(kernel) -> str | None:
+def _infer_kernels_flash_variant(kernel: object) -> str | None:
     if hasattr(kernel, "fwd") and hasattr(kernel, "varlen_fwd"):
         return "flash_attn2"
     if hasattr(kernel, "flash_attn_func") and hasattr(kernel, "flash_attn_varlen_func"):
@@ -273,6 +284,40 @@ def _ensure_flash_kernels_loaded(implementation: str) -> tuple[object, str]:
     return loaded
 
 
+def flash_kernel_unsupported_reason(implementation: str) -> str | None:
+    """Return why this platform or GPU cannot run a manifest-locked kernel, or None if it can.
+
+    The answer comes from kernels.lock, the manifest, and the current CUDA device, so nothing
+    is downloaded or imported. A caller choosing between FlashAttention versions may skip a
+    kernel for this reason alone; any other loading failure must raise.
+    """
+    from fastplms.registry import get_model_registry
+
+    kernel_spec = get_model_registry().attention_kernels[implementation]
+    pinned = f"{kernel_spec.repository}@{kernel_spec.revision}"
+    if locked_variant_for_this_system(kernel_spec.repository, kernel_spec.revision) is None:
+        return f"kernels.lock pins no build of {pinned} for this platform."
+
+    # `kernels` matches a build to the PyTorch backend, so on a PyTorch without CUDA the
+    # build found above is a CPU or XPU one and needs no GPU.
+    if torch.version.cuda is None:
+        return None
+
+    if not torch.cuda.is_available():
+        return f"{pinned} runs on a CUDA device, and none is visible."
+
+    capability = torch.cuda.get_device_capability()
+    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 (
+            f"{pinned} requires CUDA compute capability {required} or newer; "
+            f"this GPU has {observed}."
+        )
+
+    return None
+
+
 def _kernels_flash_forward(
     query_states: torch.Tensor,
     key_states: torch.Tensor,
@@ -371,7 +416,7 @@ def _kernels_flash_varlen_forward(
 # before the kernel call and restore the original padded batch shape afterward.
 class IndexFirstAxis(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, input, indices) -> torch.Tensor:
+    def forward(ctx: Any, input: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
         # input: (n, ...); indices: (m,)
         ctx.save_for_backward(indices)
         if input.ndim < 2:
@@ -391,7 +436,7 @@ class IndexFirstAxis(torch.autograd.Function):
         ).reshape(-1, *other_shape)
 
     @staticmethod
-    def backward(ctx, grad_output) -> tuple[torch.Tensor, None]:
+    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]:
         # grad_output: (m, ...)
         (indices,) = ctx.saved_tensors
         if grad_output.ndim < 2:
@@ -412,7 +457,9 @@ class IndexFirstAxis(torch.autograd.Function):
 
 class IndexPutFirstAxis(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, values, indices, first_axis_dim) -> torch.Tensor:
+    def forward(
+        ctx: Any, values: torch.Tensor, indices: torch.Tensor, first_axis_dim: int
+    ) -> torch.Tensor:
         # values: (m, ...); indices: (m,)
         ctx.save_for_backward(indices)
         if indices.ndim != 1:
@@ -432,7 +479,7 @@ class IndexPutFirstAxis(torch.autograd.Function):
         return output  # (n, ...)
 
     @staticmethod
-    def backward(ctx, grad_output) -> tuple[torch.Tensor, None, None]:
+    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:
         # grad_output: (n, ...)
         (indices,) = ctx.saved_tensors
         return grad_output[indices], None, None  # (m, ...), None, None
@@ -451,7 +498,7 @@ def _select_first_axis(states: torch.Tensor, indices: torch.Tensor) -> torch.Ten
     # states: (n, ...); indices: (m,)
     if states.requires_grad:
         selected: torch.Tensor = index_first_axis(states, indices)  # (m, ...)
-        return selected
+        return selected  # (m, ...)
     return states[indices]  # (m, ...)
 
 
@@ -528,7 +575,7 @@ def _unpad_input(
         indices,
         (cu_seqlens, cu_seqlens),
         (max_seqlen, max_seqlen),
-    )
+    )  # (t, h, d), (t, h, d), (t, h, d), (t,), ((b + 1,), (b + 1,)), (int, int)
 
 
 def _validate_flash_padding_mask(
@@ -786,7 +833,7 @@ def resolve_attention_backend(
     return resolved
 
 
-def get_attn_implementation(config) -> str:
+def get_attn_implementation(config: PretrainedConfig) -> str:
     """Read the Transformers attention setting, defaulting to SDPA."""
     requested = getattr(config, "_attn_implementation", None)
     if requested is None:
@@ -794,7 +841,7 @@ def get_attn_implementation(config) -> str:
     return resolve_attention_backend(requested).value
 
 
-def set_config_attn_implementation(config, implementation: str) -> str:
+def set_config_attn_implementation(config: PretrainedConfig, implementation: str) -> str:
     """Set both the Transformers field and the internal dispatch field."""
     resolved = resolve_attention_backend(implementation).value
     if hasattr(config, "_attn_implementation_internal"):
@@ -823,7 +870,7 @@ def get_attention_mask(
     """
     # attention_mask: (b, l) or None
     if attention_mask is None:
-        return None, None, None
+        return None, None, None  # (None, None, None)
 
     if attention_mask.ndim != 2:
         raise ValueError(
@@ -850,12 +897,15 @@ def get_attention_mask(
             raise RuntimeError(
                 "'flex_attention' was requested, but torch.create_block_mask is unavailable."
             )
-        def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
+        def mask_mod(
+            batch_idx: torch.Tensor, head_idx: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor
+        ) -> torch.Tensor:
+            # batch_idx, head_idx, q_idx, kv_idx: ()
             del head_idx, q_idx
             # Match eager and SDPA: padding masks suppress invalid keys only.
             # Invalid queries still attend to real keys and therefore remain
             # finite; downstream residue masks exclude their outputs.
-            return attention_mask_2d[batch_idx, kv_idx]
+            return attention_mask_2d[batch_idx, kv_idx]  # ()
 
         flex_block_mask = _get_flex_block_mask(
             mask_pattern=attention_mask_2d,

fastplms/attention/_kernel_lock.py

@@ -2,13 +2,26 @@
 
 from __future__ import annotations
 
+import hashlib
 import json
 import os
+import re
 
 from pathlib import Path
 from typing import Any
 
 
+# kernels.lock records, per locked build variant, one SHA-256 over the variant's relative
+# file paths and the Git or Git LFS object IDs of their contents. kernels 0.17 no longer
+# writes or checks these digests, so this module checks them before anything is imported.
+_VARIANT_HASH_TYPE = "git_lfs_concat"
+_VARIANT_DIGEST = re.compile(r"sha256-[0-9a-f]{64}")
+# The Hub cache names a Git blob by its 40-character SHA-1 object ID and Git LFS content
+# by its 64-character SHA-256.
+_GIT_OBJECT_ID_LENGTH = 40
+_LFS_OBJECT_ID_LENGTH = 64
+
+
 def require_kernels_package() -> None:
     """Fail early when the precompiled-kernel runtime is not installed."""
     try:
@@ -51,6 +64,68 @@ def _locked_entry(lock_path: Path, repository: str) -> dict[str, Any]:
     return matches[0]
 
 
+def _locked_variant_digests(entry: dict[str, Any]) -> dict[str, str]:
+    """Each locked build variant of one kernels.lock entry, mapped to its digest."""
+    variants = entry.get("variants")
+    if not isinstance(variants, dict) or not variants:
+        raise RuntimeError(f"kernels.lock locks no build variants for {entry.get('repo_id')!r}.")
+    digests: dict[str, str] = {}
+    for variant_name, variant_lock in variants.items():
+        expected_hash = variant_lock.get("hash") if isinstance(variant_lock, dict) else None
+        if (
+            not isinstance(expected_hash, str)
+            or _VARIANT_DIGEST.fullmatch(expected_hash) is None
+            or variant_lock.get("hash_type") != _VARIANT_HASH_TYPE
+        ):
+            raise RuntimeError(f"The kernel lock for {variant_name} has no valid SHA-256 digest.")
+        digests[variant_name] = expected_hash
+    return digests
+
+
+def _git_blob_object_id(contents: bytes) -> bytes:
+    """Return the SHA-1 object ID Git assigns to a blob with these contents."""
+    return hashlib.sha1(b"blob %d\0" % len(contents) + contents).digest()
+
+
+def validate_variant_digest(snapshot: Path, variant_name: str, expected_hash: str) -> None:
+    """Check one cached build variant against its kernels.lock digest before import.
+
+    Snapshot files link into the Hub cache's content-addressed blobs. Each linked file
+    contributes its path relative to the variant, then the object ID its contents hash to:
+    a Git blob ID when the cache names the blob by SHA-1, a Git LFS SHA-256 otherwise.
+    Files that are not links are skipped, because importing a kernel writes bytecode
+    beside it.
+    """
+    variant_root = snapshot / "build" / variant_name
+    linked_files: list[tuple[bytes, Path]] = []
+    for directory, _, file_names in os.walk(variant_root):
+        for file_name in file_names:
+            path = Path(directory) / file_name
+            if path.is_symlink():
+                relative_name = path.relative_to(variant_root).as_posix().encode("utf-8")
+                linked_files.append((relative_name, path))
+
+    digest = hashlib.sha256()
+    for relative_name, path in sorted(linked_files):
+        contents = path.read_bytes()
+        object_id_length = len(path.resolve().name)
+        if object_id_length == _GIT_OBJECT_ID_LENGTH:
+            object_id = _git_blob_object_id(contents)
+        elif object_id_length == _LFS_OBJECT_ID_LENGTH:
+            object_id = hashlib.sha256(contents).digest()
+        else:
+            raise RuntimeError(f"Unexpected Hub cache blob name behind {path}.")
+        digest.update(relative_name)
+        digest.update(object_id)
+
+    received_hash = f"sha256-{digest.hexdigest()}"
+    if received_hash != expected_hash:
+        raise RuntimeError(
+            f"The cached kernel variant {variant_name} hashes to {received_hash}, but "
+            f"kernels.lock records {expected_hash}."
+        )
+
+
 def _offline_mode() -> bool:
     """Return whether Hub access was explicitly disabled for this process."""
 
@@ -78,7 +153,7 @@ def _offline_snapshot_path(repository: str, revision: str) -> Path:
     if not snapshot.is_dir():
         raise RuntimeError(
             f"The exact offline kernel snapshot {repository}@{revision} is not cached under "
-            f"{cache_root}. Run `kernels download` before enabling offline mode."
+            f"{cache_root}. Load the kernel once with Hub access before enabling offline mode."
         )
     if repository_root not in snapshot.resolve().parents:
         raise RuntimeError(f"Refusing kernel snapshot outside its cache repository: {snapshot}")
@@ -88,7 +163,7 @@ def _offline_snapshot_path(repository: str, revision: str) -> Path:
 def _load_offline_locked_kernel(
     repository: str,
     revision: str,
-    variant_locks: dict[str, object],
+    variant_digests: dict[str, str],
 ) -> object:
     """Validate and import the one compatible variant from a sparse Hub snapshot."""
     snapshot = _offline_snapshot_path(repository, revision)
@@ -97,7 +172,7 @@ def _load_offline_locked_kernel(
         raise RuntimeError(f"The cached kernel snapshot has no build directory: {snapshot}")
 
     cached_names = sorted(entry.name for entry in build_root.iterdir() if entry.is_dir())
-    unexpected = sorted(set(cached_names).difference(variant_locks))
+    unexpected = sorted(set(cached_names).difference(variant_digests))
     if unexpected:
         raise RuntimeError(
             f"The cached {repository}@{revision} snapshot contains unlocked variants: "
@@ -106,7 +181,6 @@ def _load_offline_locked_kernel(
 
     try:
         from kernels import get_local_kernel
-        from kernels.utils import validate_kernel
         from kernels.variants import get_variants_local, resolve_variants
     except ImportError as error:
         raise RuntimeError(
@@ -130,47 +204,91 @@ def _load_offline_locked_kernel(
             f"found {names}."
         )
     variant_name = compatible[0].variant_str
-    variant_lock = variant_locks.get(variant_name)
-    expected_hash = getattr(variant_lock, "hash", None)
-    if not isinstance(expected_hash, str) or not expected_hash.startswith("sha256-"):
-        raise RuntimeError(f"The kernel lock for {variant_name} has no valid SHA-256 digest.")
-
-    # Hash validation deliberately happens before import. This operates on the
-    # sparse snapshot produced by `kernels download` and avoids Hub 1.23's
-    # full-snapshot completeness check in offline mode.
-    validate_kernel(repo_path=snapshot, variant=variant_name, hash=expected_hash)
+
+    # Hash validation deliberately happens before import. It reads the sparse snapshot
+    # directly, so no Hub API judges whether the snapshot is complete.
+    validate_variant_digest(snapshot, variant_name, variant_digests[variant_name])
     return get_local_kernel(build_root / variant_name)
 
 
-def load_locked_kernel(repository: str, revision: str) -> object:
-    """Download, hash-validate, then import one immutable precompiled kernel."""
-    require_kernels_package()
+def _compatible_locked_variants(repository: str, variant_names: list[str]) -> list[str]:
+    """Return the locked build variants `kernels` can load on this system, preferred first."""
     try:
-        from kernels import get_local_kernel, install_kernel
-        from kernels.lockfile import KernelLock
+        from kernels.variants import parse_variant, resolve_variants
     except ImportError as error:
         raise RuntimeError(
             "Precompiled FlashAttention requires requirements/features/flash.in."
         ) from error
 
-    lock_path = _kernel_lock_path()
-    kernel_lock = KernelLock.from_json(_locked_entry(lock_path, repository))
-    if kernel_lock.sha != revision:
+    try:
+        locked_variants = [parse_variant(variant_name) for variant_name in variant_names]
+    except ValueError as error:
+        raise RuntimeError(f"kernels.lock contains an invalid {repository} variant.") from error
+    compatible, _ = resolve_variants(locked_variants)
+    return [variant.variant_str for variant in compatible]
+
+
+def _preferred_locked_variant(repository: str, revision: str, variant_names: list[str]) -> str:
+    """Return the locked build variant `kernels` prefers on this system."""
+    compatible = _compatible_locked_variants(repository, variant_names)
+    if not compatible:
+        raise RuntimeError(
+            f"kernels.lock locks no build of {repository}@{revision} for this system; "
+            f"locked variants: {', '.join(sorted(variant_names))}."
+        )
+    return compatible[0]
+
+
+def _pinned_variant_digests(repository: str, revision: str) -> dict[str, str]:
+    """Return the locked build digests of a kernel whose kernels.lock entry pins `revision`."""
+    entry = _locked_entry(_kernel_lock_path(), repository)
+    locked_revision = entry.get("sha")
+    if locked_revision != revision:
         raise RuntimeError(
             f"The typed manifest pins {repository}@{revision}, but kernels.lock pins "
-            f"{kernel_lock.sha}."
+            f"{locked_revision}."
         )
+    return _locked_variant_digests(entry)
+
+
+def locked_variant_for_this_system(repository: str, revision: str) -> str | None:
+    """Return the locked build `kernels` would load here, or None when kernels.lock pins none.
+
+    Only kernels.lock is read, so a caller can tell a platform the lock does not cover
+    apart from a download, digest, or import failure.
+    """
+    variant_digests = _pinned_variant_digests(repository, revision)
+    compatible = _compatible_locked_variants(repository, list(variant_digests))
+    return compatible[0] if compatible else None
+
+
+def load_locked_kernel(repository: str, revision: str) -> object:
+    """Download, hash-validate, then import one immutable precompiled kernel."""
+    require_kernels_package()
+    try:
+        from huggingface_hub import snapshot_download
+        from kernels import get_local_kernel
+    except ImportError as error:
+        raise RuntimeError(
+            "Precompiled FlashAttention requires requirements/features/flash.in."
+        ) from error
+
+    variant_digests = _pinned_variant_digests(repository, revision)
 
     if _offline_mode():
-        return _load_offline_locked_kernel(repository, revision, kernel_lock.variants)
-
-    # `install_kernel` downloads data without importing it and validates the
-    # selected build against the tracked variant hash. Only then is the exact
-    # validated path imported directly. Offline mode uses the sparse-cache
-    # resolver above because Hub 1.23 rejects partial snapshots as incomplete.
-    validated_path = install_kernel(
-        repository,
-        revision=kernel_lock.sha,
-        variant_locks=kernel_lock.variants,
+        return _load_offline_locked_kernel(repository, revision, variant_digests)
+
+    # Only a locked build can be selected, and only its files are downloaded, at the
+    # immutable revision. The download imports nothing; the digest check runs first.
+    variant_name = _preferred_locked_variant(repository, revision, list(variant_digests))
+    snapshot = Path(
+        snapshot_download(
+            repository,
+            repo_type="kernel",
+            revision=revision,
+            allow_patterns=[f"build/{variant_name}/*"],
+            cache_dir=os.environ.get("KERNELS_CACHE") or None,
+        )
     )
-    return get_local_kernel(validated_path)
+    validate_variant_digest(snapshot, variant_name, variant_digests[variant_name])
+    return get_local_kernel(snapshot / "build" / variant_name)

fastplms/attention/interfaces.py

@@ -7,7 +7,7 @@ import torch
 from collections.abc import Mapping
 from functools import partial
 from typing import Any
-from transformers import AttentionInterface, AttentionMaskInterface
+from transformers import AttentionInterface, AttentionMaskInterface, PretrainedConfig
 
 from ._auto import (
     AUTO_ATTENTION,
@@ -94,7 +94,7 @@ class FastPLMsAttentionMixin:
 
     _supports_sdpa = True
     _supports_flex_attn = True
-    # Transformers 5.13 uses the singular flag during model construction. A
+    # Transformers checks the singular flag during model construction. A
     # family opts in only when its manifest entry advertises at least one of
     # the two FastPLMs kernels-only FlashAttention implementations.
     _supports_flash_attn = False
@@ -161,7 +161,7 @@ class FastPLMsAttentionMixin:
             allow_all_kernels=False,
         )
 
-    def __init__(self, config, *args: Any, **kwargs: Any) -> None:
+    def __init__(self, config: PretrainedConfig, *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
@@ -355,7 +355,7 @@ def _resolve_auto_attention_before_forward(
 def validate_transformers_attention_interfaces() -> None:
     """Verify that Transformers exposes functions and masks for every backend.
 
-    Transformers 5.13 registers these canonical names. The FastPLMs function
+    The validated Transformers registers these canonical names. The FastPLMs function
     overrides remain instance-local and do not replace process-global handlers.
     """
     function_registry = FASTPLMS_ATTENTION_FUNCTIONS

fastplms/digests.py

@@ -0,0 +1,35 @@
+"""SHA-256 digests of files and JSON values, the identities FastPLMs records and compares.
+
+This file exists twice, byte for byte: here and as ``features/digests.py``. ``features`` loads as a standalone
+package (``foundry.embedding.private_store``) and imports nothing outside itself, so it cannot reach this
+module. ``tests/tier1_unit/test_features_package_is_self_contained.py`` fails when the two differ.
+"""
+
+from __future__ import annotations
+
+import hashlib
+
+from pathlib import Path
+from typing import Any
+
+from .json_files import compact_json
+
+
+FILE_READ_BYTES = 1024 * 1024
+
+
+def file_sha256(path: str | Path) -> str:
+    """Return the SHA-256 of a file's bytes, read in blocks so a checkpoint never sits in memory."""
+
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as handle:
+        while block := handle.read(FILE_READ_BYTES):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def json_sha256(value: Any, *, ensure_ascii: bool = True, allow_nan: bool = True) -> str:
+    """Return the SHA-256 of ``value`` in its compact, key-sorted JSON form (``compact_json``)."""
+
+    encoded = compact_json(value, ensure_ascii=ensure_ascii, allow_nan=allow_nan).encode("utf-8")
+    return hashlib.sha256(encoded).hexdigest()

fastplms/embeddings/__init__.py

@@ -1,6 +1,9 @@
 """Ordered, residue-aware protein embedding utilities."""
 
-from .pooling import POOLING_NAMES, Pooler, pagerank_weights
+from .feature_runs import embed_into_features
+from .pooling import (
+    POOLING_NAMES, POOLING_SEMANTICS_TOKENS, TOKEN_POOLING_NAMES, Pooler, pagerank_weights, pool_token_rows,
+)
 from .runner import (
     EmbeddingMixin,
     embed_dataset,
@@ -24,12 +27,21 @@ from .storage import (
     tensor_sha256,
     update_sqlite_run_metadata,
 )
+from .taps import (
+    HiddenTap, LayerAccumulator, ReducedTap, RowSelection, SparseResidueTap, StreamingTap, TapBatch,
+)
+from .token_batches import BatchGeometry, TokenTapExecutor, plan_geometry_batches, plan_token_batches
+from .token_runs import embed_token_features
+from .tokens import ResidueVocabulary
 from .types import (
     EmbeddingBatch,
     EmbeddingInput,
     EmbeddingRecord,
     EmbeddingResult,
     LazyTensorReference,
+    TapRecord,
+    TapResult,
+    TapRunReceipt,
     TensorValue,
 )
 
@@ -37,17 +49,34 @@ from .types import (
 __all__ = [
     "DEFAULT_SHARD_SIZE",
     "POOLING_NAMES",
+    "POOLING_SEMANTICS_TOKENS",
+    "TOKEN_POOLING_NAMES",
+    "BatchGeometry",
     "EmbeddingBatch",
     "EmbeddingInput",
     "EmbeddingMixin",
     "EmbeddingRecord",
     "EmbeddingResult",
+    "HiddenTap",
+    "LayerAccumulator",
     "LazyTensorReference",
     "Pooler",
+    "ReducedTap",
+    "ResidueVocabulary",
+    "RowSelection",
+    "SparseResidueTap",
+    "StreamingTap",
+    "TapBatch",
+    "TapRecord",
+    "TapResult",
+    "TapRunReceipt",
     "TensorValue",
+    "TokenTapExecutor",
     "append_sqlite_records",
     "convert_legacy_sqlite",
     "embed_dataset",
+    "embed_into_features",
+    "embed_token_features",
     "garbage_collect_safetensors_generations",
     "initialize_sqlite_run",
     "iter_fasta",
@@ -57,6 +86,9 @@ __all__ = [
     "load_sqlite_result",
     "pagerank_weights",
     "parse_fasta",
+    "plan_geometry_batches",
+    "plan_token_batches",
+    "pool_token_rows",
     "save_result",
     "save_safetensors_result",
     "save_sqlite_result",

fastplms/embeddings/batches.py

@@ -13,7 +13,9 @@ from torch import Tensor
 from .identity import _model_device
 from .inputs import _planned_batches
 from .pooling import Pooler
-from .types import EmbeddingBatch, EmbeddingInput, EmbeddingRecord
+from .taps import HiddenTap, ReducedTap, SparseResidueTap, StreamingTap, TapBatch, TapPlan
+from .types import EmbeddingBatch, EmbeddingInput, EmbeddingRecord, TapRecord
+from ..features.layouts import TopKRow, validate_topk
 
 
 _MAX_PARTI_RESIDUES = 2_048
@@ -36,7 +38,7 @@ def select_hidden_state_embeddings(
     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)
+    # last_hidden_state: (b, l, d); hidden_states: (b, l, d) per entry
     if store_all_hidden_states:
         if not hidden_states:
             raise ValueError("store_all_hidden_states requires model hidden states.")
@@ -101,6 +103,61 @@ def _biological_residue_mask(
     return M  # (b, l)
 
 
+def canonical_residue_ids(sequence: str, tokenizer: Any) -> list[int]:
+    """Validate an already normalized protein's one-token-per-residue representation.
+
+    This does not normalize input or create sequence identity. Canonical callers supply their
+    verified inventory text; unknown residues, including J in ESMC, fail before inference.
+    """
+    if not sequence or not sequence.isascii() or not sequence.isalpha() or not sequence.isupper():
+        raise ValueError("Canonical feature input must be an already normalized uppercase protein.")
+    ids = tokenizer.convert_tokens_to_ids(list(sequence))
+    special = set(tokenizer.all_special_ids)
+    if (not isinstance(ids, list) or len(ids) != len(sequence)
+            or any(type(token) is not int or token in special for token in ids)):
+        raise ValueError("A canonical residue has no non-special tokenizer representation.")
+    return ids
+
+
+def _tokenized_batch(
+    model: Any,
+    sequences: list[str],
+    *,
+    tokenizer: Any,
+    max_length: int | None,
+    truncate: bool,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Tokenize one batch on the model device: input IDs, attention mask, and residue mask M."""
+
+    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)
+    return input_ids, attention_mask, M  # each: (b, l)
+
+
 def _generic_embedding_batch(
     model: Any,
     sequences: list[str],
@@ -140,32 +197,13 @@ def _generic_embedding_batch(
     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)
+    input_ids, attention_mask, M = _tokenized_batch(
+        model,
+        sequences,
+        tokenizer=tokenizer,
+        max_length=max_length,
+        truncate=truncate,
+    )  # each: (b, l)
     if need_attentions:
         # Validate l before either the backbone or its quadratic attention graph
         # is materialized. M has shape (b, l).
@@ -189,6 +227,24 @@ def _generic_embedding_batch(
     )
 
 
+def _sparse_residue_rows(tap: SparseResidueTap, batch: TapBatch) -> list[TopKRow]:
+    """Validate and own each sequence's sparse biological-residue outputs on the CPU."""
+    rows = tap.reduce(batch)
+    if not isinstance(rows, Sequence) or len(rows) != batch.X.shape[0]:
+        raise ValueError("A sparse residue reducer must return one TopKRow per sequence.")
+    output = []
+    for row, length in zip(rows, batch.residue_mask.sum(dim=1).tolist(), strict=True):
+        if not isinstance(row, TopKRow):
+            raise TypeError("A sparse residue reducer must return TopKRow values.")
+        validate_topk(row.indices, row.values, tap.codebook_size, tap.sparse_count)
+        if row.values.shape[0] != length:
+            raise ValueError("Sparse residue output must retain every biological residue in order.")
+        output.append(TopKRow(
+            row.indices.detach().cpu().clone(), row.values.detach().cpu().clone(),
+        ))
+    return output
+
+
 @dataclass(eq=False)
 class BatchExecutor:
     """Model and batch policy for one bounded embedding window at a time."""
@@ -312,7 +368,7 @@ class BatchExecutor:
             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 = (
+            valid_embedding_shape = (
                 X.ndim == 3
                 and X.shape[0] == len(batch_records)
                 and X.shape[-1] > 0
@@ -327,7 +383,7 @@ class BatchExecutor:
                 and X.shape[-1] > 0
                 and M.shape == (X.shape[0], X.shape[2])
             )
-            if not (valid_X_shape or valid_all_states_shape):
+            if not (valid_embedding_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 "
@@ -335,7 +391,7 @@ class BatchExecutor:
                 )
             if not bool(M.any(dim=1).all()):
                 raise ValueError("Every embedding sample must contain a biological residue.")
-            finite_selected = (  # X.shape
+            finite_selected = (  # (b, l, d) or (b, n_states, l, d)
                 torch.isfinite(X) | ~M.unsqueeze(-1)
                 if X.ndim == 3
                 else torch.isfinite(X) | ~M[:, None, :, None]
@@ -347,6 +403,11 @@ class BatchExecutor:
                 _validate_parti_length(M)
             if self.dtype is not None:
                 X = X.to(dtype=self.dtype)  # unchanged shape
+                selected = M.unsqueeze(-1) if X.ndim == 3 else M[:, None, :, None]  # (b, l, 1) or (b, 1, l, 1)
+                if not bool((torch.isfinite(X) | ~selected).all()):
+                    raise ValueError(
+                        "Embedding dtype conversion produced non-finite biological residues."
+                    )
 
             if self.full_embeddings:
                 if X.ndim == 4:
@@ -377,3 +438,224 @@ class BatchExecutor:
             for position in range(window_start, window_start + len(window_records))
         ]
         return new_records, pool_slices
+
+
+def _tap_states(
+    model: Any,
+    sequences: list[str],
+    *,
+    tokenizer: Any | None,
+    max_length: int | None,
+    truncate: bool,
+    layers: tuple[int, ...],
+    streaming: tuple[StreamingTap, ...] = (),
+    require_residue_identity: bool = False,
+) -> tuple[dict[int, Tensor], Tensor, Tensor, dict[str, Tensor]]:
+    """Record ``layers`` in one forward pass; return the states, the token mask, and M."""
+
+    resolved_tokenizer = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None)
+    if resolved_tokenizer is None:
+        raise ValueError("A tokenizer is required for this model's embedding path.")
+    expected_ids = (
+        [canonical_residue_ids(sequence, resolved_tokenizer) for sequence in sequences]
+        if require_residue_identity else None
+    )
+    input_ids, attention_mask, M = _tokenized_batch(
+        model,
+        sequences,
+        tokenizer=resolved_tokenizer,
+        max_length=max_length,
+        truncate=truncate,
+    )  # each: (b, l)
+    token_mask = attention_mask.to(dtype=torch.bool)  # (b, l)
+    if expected_ids is not None:
+        for index, expected in enumerate(expected_ids):
+            # The actual biological mask must select each original residue exactly once,
+            # in order. This checks tokenization, padding and cropping before the encoder.
+            if input_ids[index, M[index]].tolist() != expected:
+                raise ValueError(
+                    "Tokenizer and biological mask do not preserve original residue positions."
+                )
+    if not bool(M.any(dim=1).all()):
+        raise ValueError("Every embedding sample must contain a biological residue.")
+    streamed: dict[str, Tensor] = {}
+    if streaming:
+        if getattr(model, "embedding_streaming_tap_support", False) is not True:
+            raise ValueError("This model does not support streaming hidden-state taps.")
+        accumulators = [(tap, tap.begin()) for tap in streaming]
+        stream_layers = tuple(sorted({layer for tap in streaming for layer in tap.layers}))
+        seen: list[int] = []
+
+        def consume(layer: int, X: Tensor) -> None:
+            # X: (b, l, d), borrowed until this callback returns.
+            if len(seen) >= len(stream_layers) or layer != stream_layers[len(seen)]:
+                raise RuntimeError("Streaming hidden states arrived out of plan order.")
+            if X.ndim != 3 or X.shape[:2] != M.shape:
+                raise ValueError("Streaming hidden states are not token-aligned.")
+            if not bool((torch.isfinite(X) | ~M.unsqueeze(-1)).all()):
+                raise ValueError("Biological residue embeddings produced non-finite output.")
+            seen.append(layer)
+            batch = TapBatch(X, token_mask, M)
+            for tap, accumulator in accumulators:
+                if layer in tap.layers:
+                    accumulator.update(layer, batch)
+
+        states = model._embed_taps(
+            input_ids, attention_mask, layers, stream_layers=stream_layers, state_consumer=consume,
+        )
+        if tuple(seen) != stream_layers:
+            raise RuntimeError("The encoder did not deliver every streaming hidden state.")
+        for tap, accumulator in accumulators:
+            Y = accumulator.finish()  # (b, l, c)
+            if (not isinstance(Y, Tensor) or Y.ndim != 3
+                    or Y.shape[:2] != M.shape or Y.shape[2] == 0):
+                raise ValueError(
+                    f"Streaming tap {tap.name!r} must return token-aligned (b, l, c) features."
+                )
+            if Y.device != M.device or not Y.is_floating_point():
+                raise ValueError("Streaming features must be floating tensors on the input device.")
+            if not bool((torch.isfinite(Y) | ~M.unsqueeze(-1)).all()):
+                raise ValueError(
+                    f"Streaming tap {tap.name!r} produced non-finite biological residues."
+                )
+            streamed[tap.name] = Y
+    else:
+        states = model._embed_taps(input_ids, attention_mask, layers)  # each: (b, l, d)
+    # states: (b, l, d) per requested layer; token_mask, M: (b, l); streamed: (b, l, c) per streaming tap
+    return states, token_mask, M, streamed  # ((b, l, d), ...), (b, l), (b, l), {tap: (b, l, c)}
+
+
+def _reduced_rows(tap: ReducedTap, batch: TapBatch) -> list[Tensor]:
+    """Apply a reducer and split its output into one CPU tensor per sequence."""
+
+    # batch.X: (b, l, d)
+    Y = tap.reduce(batch)  # (b, ...)
+    if not isinstance(Y, Tensor):
+        raise TypeError(f"The reducer of tap {tap.name!r} must return a Tensor.")
+    if Y.ndim == 0 or Y.shape[0] != batch.X.shape[0]:
+        raise ValueError(
+            f"The reducer of tap {tap.name!r} must return one row per sequence, shape "
+            f"(b, ...) with b={batch.X.shape[0]}; it returned {tuple(Y.shape)}."
+        )
+    if Y.is_floating_point() and not bool(torch.isfinite(Y).all()):
+        raise ValueError(f"The reducer of tap {tap.name!r} produced non-finite output.")
+    return list(Y.detach().cpu().unbind(0))  # b tensors of shape (...), Y without its batch axis
+
+
+@dataclass(eq=False)
+class TapExecutor:
+    """Model and batch policy for one bounded window of a tap plan.
+
+    Each batch runs one forward pass that records every tapped state and stops after the
+    deepest. A hidden tap's dtype overrides the run dtype. Conversion starts from the original
+    state for each tap, so a low-precision residue output cannot quantize a pooled sibling.
+    A reducer receives the run dtype and its output keeps the dtype it returns.
+    """
+
+    model: Any
+    plan: TapPlan
+    batch_size: int
+    max_tokens_per_batch: int | None
+    max_length: int | None
+    truncate: bool
+    tokenizer: Any | None
+    dtype: torch.dtype | None
+    attention_backend: str | None
+    require_residue_identity: bool = False
+    poolers: dict[str, Pooler] = field(init=False)
+
+    def __post_init__(self) -> None:
+        pooled_streams = [tap.name for tap in self.plan.taps if isinstance(tap, StreamingTap) and tap.pooling is not None]
+        if pooled_streams:
+            raise ValueError(f"Pooled streaming taps {pooled_streams} run in a token run (TokenTapExecutor) only.")
+        self.poolers = {
+            tap.name: Pooler(tap.pooling)
+            for tap in self.plan.taps
+            if isinstance(tap, HiddenTap) and tap.pooling is not None
+        }
+
+    def run_window(
+        self,
+        window_records: Sequence[EmbeddingInput],
+        *,
+        window_start: int,
+    ) -> tuple[list[TapRecord], dict[str, dict[str, tuple[int, int]]]]:
+        """Restore source order after length-bucketed inference; return each pooled tap's slices."""
+
+        pool_slices: dict[str, dict[str, tuple[int, int]]] = {}
+        window_results: dict[int, TapRecord] = {}
+        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_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
+            ]
+            states, token_mask, M, streamed = _tap_states(  # states: each (b, l, d); masks: (b, l)
+                self.model,
+                sequences,
+                tokenizer=self.tokenizer,
+                max_length=self.max_length,
+                truncate=self.truncate,
+                layers=self.plan.captured_layers,
+                streaming=tuple(tap for tap in self.plan.taps if isinstance(tap, StreamingTap)),
+                require_residue_identity=self.require_residue_identity,
+            )
+            if not bool(M.any(dim=1).all()):
+                raise ValueError("Every embedding sample must contain a biological residue.")
+            for X in states.values():  # each: (b, l, d)
+                if not bool((torch.isfinite(X) | ~M.unsqueeze(-1)).all()):
+                    raise ValueError("Biological residue embeddings produced non-finite output.")
+            outputs: dict[str, list[Tensor] | list[TopKRow]] = {}
+            for tap, layer in zip(self.plan.taps, self.plan.layers, strict=True):
+                if isinstance(tap, StreamingTap):
+                    outputs[tap.name] = _residue_embeddings(streamed[tap.name], M)  # each: (r_i, c)
+                    continue
+                X = states[layer]  # (b, l, d)
+                dtype = self.dtype
+                if isinstance(tap, HiddenTap) and tap.dtype is not None:
+                    dtype = tap.dtype
+                if dtype is not None:
+                    X = X.to(dtype=dtype)  # (b, l, d), from the original captured state
+                    if not bool((torch.isfinite(X) | ~M.unsqueeze(-1)).all()):
+                        raise ValueError(
+                            "Tap dtype conversion produced non-finite biological residues."
+                        )
+                if isinstance(tap, SparseResidueTap):
+                    batch = TapBatch(X=X, token_mask=token_mask, residue_mask=M)
+                    outputs[tap.name] = _sparse_residue_rows(tap, batch)  # each pair: (r_i,k)
+                elif isinstance(tap, ReducedTap):
+                    batch = TapBatch(X=X, token_mask=token_mask, residue_mask=M)
+                    outputs[tap.name] = _reduced_rows(tap, batch)  # each: (...)
+                elif tap.pooling is None:
+                    outputs[tap.name] = _residue_embeddings(X, M)  # each: (r_i, d)
+                else:
+                    pooler = self.poolers[tap.name]
+                    Y = pooler(X, M, attention_backend=self.attention_backend)  # (b, n_poolers * d)
+                    pool_slices[tap.name] = pooler.output_slices(X.shape[-1])
+                    outputs[tap.name] = list(Y.detach().cpu().unbind(0))  # each: (n_poolers * d,)
+            for offset, (position, record) in enumerate(
+                zip(local_positions, batch_records, strict=True)
+            ):
+                window_results[window_start + position] = TapRecord(
+                    record.id,
+                    record.sequence,
+                    {name: values[offset] for name, values in outputs.items()},
+                    # This correspondence was proven against actual token IDs and M before
+                    # inference, not inferred from a tensor's row count.
+                    tuple(range(len(sequences[offset]))) if self.require_residue_identity else None,
+                )
+
+        new_records = [
+            window_results[position]
+            for position in range(window_start, window_start + len(window_records))
+        ]
+        return new_records, pool_slices

fastplms/embeddings/feature_runs.py

@@ -0,0 +1,232 @@
+"""Run a tap plan and persist each tap into its own feature store.
+
+This is the one place a model runs to fill the store. ``embed_into_features`` embeds only the
+sequences the stores lack, takes one forward pass per batch for every tap, and writes each tap's
+rows into the store of its key. A second call with the same sequences runs no model at all.
+
+Each tap becomes one feature, so a run that taps the last hidden state, a mean-pooled vector, and
+max-pooled sparse-autoencoder codes fills three stores from one pass. The store's layout decides
+how a tap's tensor is stored: a pooled vector goes in dense, per-residue rows go in ragged, and a
+sparse-autoencoder vector goes in csr, compressed to its exactly non-zero codes.
+
+The caller owns the keys, because the key composes the model, its revision, the autoencoder, the
+layer, the pooling, the dtype, and the residue limit, and only the caller knows the pinned
+revisions. ``foundry.embedding.feature_key`` computes them.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from contextlib import ExitStack
+from pathlib import Path
+from typing import Any, Protocol
+from torch import Tensor
+
+from .runner import embed_dataset
+from .taps import Tap
+from .types import TapRecord, TapRunReceipt
+from ..features.layouts import DENSE, RAGGED, RAGGED_TOPK, SparseRow, TopKRow
+from ..features.store import FeatureStore, SegmentReceipt, SegmentWriter, StoredFeature
+
+
+class FeatureRunContract(Protocol):
+    """Scientific identity policy supplied by the caller, independent of the storage format."""
+
+    def validate(
+        self, model: Any, sequences: Sequence[str], features: Mapping[str, StoredFeature],
+        taps: Sequence[Tap], options: Mapping[str, Any],
+    ) -> None: ...
+
+    def validate_cached(self, name: str, store: FeatureStore, sequences: Sequence[str]) -> None: ...
+
+    def bind_rows(self, name: str, records: Sequence[TapRecord]) -> Sequence[Mapping[str, Any]]: ...
+
+    def before_commit(self) -> None: ...
+
+
+def embed_into_features(
+    model: Any,
+    sequences: Sequence[str],
+    root: str | Path,
+    features: Mapping[str, StoredFeature],
+    *,
+    taps: Sequence[Tap],
+    metadata: Mapping[str, Any] | None = None,
+    contract: FeatureRunContract | None = None,
+    max_part_bytes: int = 256 * 1024**2,
+    keep_special_tokens: bool = False,
+    **embed_kwargs: Any,
+) -> dict[str, SegmentReceipt]:
+    """Fill each named feature under ``root`` from one pass over the sequences it lacks.
+
+    ``features`` maps a tap name to the spec of the feature it fills, and must name every tap.
+    ``metadata`` is recorded on every segment this run commits, beside the run fingerprint.
+    ``max_part_bytes`` bounds each part's encoded tensor payload, excluding its small safetensors
+    header and row-identity sidecar. A single row exceeding it fails without being committed.
+    Remaining keyword arguments go to ``embed_dataset``. Tensor outputs are streamed one bounded
+    batch window at a time; input identities and part metadata still scale with inventory size.
+
+    A contract that keeps the special tokens (``contract.keep_special_tokens`` is True) selects the canonical
+    token path by itself, so a caller cannot forget the argument; a contract that keeps residues only (False)
+    refuses ``keep_special_tokens=True``.
+
+    ``keep_special_tokens=True`` is the canonical token path: every per-token stream holds l + 2 rows per
+    protein (row 0 CLS, rows 1..l residues, row l + 1 EOS, shape (n, d) with n = sum(l_i + 2) overall), every
+    pooled stream covers those same rows, and the run goes through ``embed_token_features`` and its
+    asynchronous writer. Its ``embed_kwargs`` are the contract's ``embedding_options`` (``max_length`` is the
+    crop in residues), and it returns the last committed segment of each stream; call ``embed_token_features``
+    for every segment.
+
+    Returns the committed segment of each feature that gained rows. A feature whose sequences were
+    all present is absent from the result, and an empty result means the model never ran.
+    """
+
+    # The contract names the stored layout: l + 2 rows (True), l rows (False), or no claim (no attribute).
+    contract_keeps = getattr(contract, "keep_special_tokens", None)
+    if contract_keeps is True:
+        keep_special_tokens = True
+    elif contract_keeps is False and keep_special_tokens:
+        raise ValueError("keep_special_tokens=True needs a contract captured with the special tokens kept.")
+
+    if keep_special_tokens:
+        from .token_batches import CANONICAL_MAX_RESIDUES
+        from .token_runs import embed_token_features
+
+        options = dict(embed_kwargs)
+        settings: dict[str, Any] = {
+            "max_residues": options.pop("max_length", CANONICAL_MAX_RESIDUES),
+            "max_sequences": options.pop("batch_size", 256),
+            "max_tokens": options.pop("max_tokens_per_batch", None) or 32768,
+            "window": options.pop("batch_window_size", 65536),
+            "dtype": options.pop("dtype", None),
+        }
+        options.pop("truncate", None)  # the canonical crop is always a prefix crop
+        if "fixed_batch_size" in options:
+            settings["fixed_batch_size"] = options.pop("fixed_batch_size")
+        if options:
+            raise ValueError(f"keep_special_tokens takes no other extraction options; received {sorted(options)}.")
+        received = embed_token_features(
+            model, sequences, root, features, taps=taps, contract=contract, metadata=metadata,
+            part_bytes=max_part_bytes, **settings,
+        )
+        return {name: group[-1] for name, group in received.items()}
+
+    if type(max_part_bytes) is not int or max_part_bytes <= 0:
+        raise ValueError("max_part_bytes must be a positive integer.")
+    if "tap_sink" in embed_kwargs:
+        raise ValueError("embed_into_features owns its tap_sink destination.")
+    tap_names = [tap.name for tap in taps]
+    if set(features) != set(tap_names):
+        raise ValueError(
+            "features must name exactly the taps this run takes.\n"
+            f"  taps:     {sorted(tap_names)}\n"
+            f"  features: {sorted(features)}"
+        )
+    for name, spec in features.items():
+        if spec.positions:
+            raise ValueError(
+                f"Feature {spec.key!r} stores the argmax residue of each code, and no tap carries "
+                f"them, so this run cannot fill it from tap {name!r}. A pipeline that computes "
+                "positions itself writes them through the store's own segment writer."
+            )
+
+    ordered = _distinct(sequences)
+    if not ordered:
+        raise ValueError("embed_into_features needs at least one sequence.")
+    if contract is None and any(
+        spec.descriptor.get("schema") == "feature_spec_v1" for spec in features.values()
+    ):
+        raise ValueError("Complete feature descriptors require a FeatureRunContract.")
+    if contract is not None:
+        contract.validate(model, ordered, features, taps, embed_kwargs)
+    stores = {name: FeatureStore.open(root, spec) for name, spec in features.items()}
+    wanted = {name: frozenset(store.missing(ordered)) for name, store in stores.items()}
+    if contract is not None:
+        for name, store in stores.items():
+            contract.validate_cached(name, store, [s for s in ordered if s not in wanted[name]])
+    to_embed = [sequence for sequence in ordered if any(sequence in group for group in wanted.values())]
+    if not to_embed:
+        return {}
+
+    receipts: dict[str, SegmentReceipt] = {}
+    with ExitStack() as stack:
+        writers: dict[str, SegmentWriter] = {}
+        # Recheck live dependencies and caller descriptors after all windows have been staged.
+        check = None if contract is None else lambda: contract.validate(
+            model, (), features, taps, embed_kwargs,
+        )
+
+        def append_window(records: Sequence[TapRecord], identity: Mapping[str, str]) -> None:
+            for name, store in stores.items():
+                embedded = [record for record in records if record.sequence in wanted[name]]
+                if not embedded:
+                    continue
+                rows = [_row_for(store.spec, record, name) for record in embedded]
+                identities = None if contract is None else contract.bind_rows(name, embedded)
+                if name not in writers:
+                    run_metadata = {
+                        **dict(metadata or {}), "input_fingerprint": identity["input_fingerprint"],
+                        "storage_policy": {"max_part_tensor_bytes": max_part_bytes},
+                    }
+                    writers[name] = stack.enter_context(store.segment(
+                        identity["run_fingerprint"], run_metadata, before_commit=check,
+                    ))
+                writers[name].append_bounded(
+                    [record.sequence for record in embedded], rows, row_metadata=identities,
+                    max_tensor_bytes=max_part_bytes,
+                )
+
+        run_receipt = embed_dataset(
+            model, to_embed, taps=list(taps), require_residue_identity=contract is not None,
+            tap_sink=append_window, **embed_kwargs,
+        )
+        if not isinstance(run_receipt, TapRunReceipt) or run_receipt.record_count != len(to_embed):
+            raise TypeError("Feature extraction did not complete delivery of every requested row.")
+        for name, writer in writers.items():
+            receipts[name] = writer.commit()
+    return receipts
+
+
+def _row_for(spec: StoredFeature, record: TapRecord, tap: str) -> Tensor | SparseRow | TopKRow:
+    """One tap's output for one sequence, in the shape its store stores."""
+
+    value = record.tensors[tap]  # (w,) pooled, (r_i, d) per residue, or a TopKRow of (r_i, k) tensors
+    if spec.layout == RAGGED_TOPK:
+        if not isinstance(value, TopKRow):
+            raise ValueError(f"Feature {spec.key!r} needs sparse TopKRow output from tap {tap!r}.")
+        return value  # TopKRow of (r_i, k) indices and values
+    if not isinstance(value, Tensor):
+        raise ValueError("Sparse residue tap output requires a ragged_topk feature.")
+    if spec.layout == DENSE:
+        if value.ndim != 1:
+            raise ValueError(
+                f"Feature {spec.key!r} is dense, so tap {tap!r} must give one vector per sequence; "
+                f"received shape {tuple(value.shape)}. A per-residue tap needs a ragged feature."
+            )
+        return value  # (w,)
+    if spec.layout == RAGGED:
+        if value.ndim != 2:
+            raise ValueError(
+                f"Feature {spec.key!r} is ragged, so tap {tap!r} must give (r_i, d) residue rows; "
+                f"received shape {tuple(value.shape)}."
+            )
+        return value  # (r_i, d)
+    if value.ndim != 1:
+        raise ValueError(
+            f"Feature {spec.key!r} is csr, so tap {tap!r} must give one vector per sequence; "
+            f"received shape {tuple(value.shape)}."
+        )
+    return SparseRow.from_dense(value)  # SparseRow of (nnz,) tensors, from a (w,) vector
+
+
+def _distinct(sequences: Sequence[str]) -> list[str]:
+    seen: set[str] = set()
+    ordered: list[str] = []
+    for sequence in sequences:
+        if sequence not in seen:
+            seen.add(sequence)
+            ordered.append(sequence)
+    return ordered
+
+
+__all__ = ["embed_into_features"]

fastplms/embeddings/identity.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import hashlib
-import json
 import platform
 import torch
 
@@ -12,12 +11,14 @@ from pathlib import Path
 from typing import Any
 from torch import Tensor
 
-from .inputs import _InputSpool
+from .pooling import POOLING_SEMANTICS
 from .storage import tensor_sha256
 from .types import EmbeddingInput
+from ..digests import json_sha256
+from ..json_files import compact_json
 
 
-_RUN_FINGERPRINT_SCHEMA_VERSION = 3
+_RUN_FINGERPRINT_SCHEMA_VERSION = 5
 _MODEL_STATE_HASH_CHUNK_BYTES = 16 * 1024**2
 
 
@@ -79,6 +80,33 @@ def _fingerprint_jsonable(value: Any) -> Any:
     }
 
 
+def _backend_tokenizer_content(backend: Any) -> str | None:
+    """The backend tokenizer's serialization, without the padding and truncation of its last call.
+
+    Transformers sets padding and truncation on the Rust tokenizer at the start of every encode
+    call, from that call's own arguments, and leaves them set. What the backend holds is
+    therefore the previous call's settings, which never change the next encoding. Hashing them
+    gave the first run in a process a different fingerprint from every identical run after it.
+    This run's own padding and truncation are fingerprinted through ``max_length``,
+    ``truncate``, and the batching policy.
+
+    The settings are cleared on a copy, so the caller's tokenizer keeps its state. A tokenizer
+    that was never called already serializes without them, so this normalization itself
+    does not change its identity. The enclosing run schema also versions pooling semantics.
+    """
+    to_str = getattr(backend, "to_str", None)
+    if not callable(to_str):
+        return None
+    serialized = to_str()
+    from_str = getattr(type(backend), "from_str", None)
+    if not callable(from_str):
+        return serialized
+    call_free = from_str(serialized)
+    call_free.no_truncation()
+    call_free.no_padding()
+    return call_free.to_str()
+
+
 def _tokenizer_content_sha256(tokenizer: Any) -> str:
     content: dict[str, Any] = {
         "init_kwargs": getattr(tokenizer, "init_kwargs", None),
@@ -93,17 +121,10 @@ def _tokenizer_content_sha256(tokenizer: Any) -> str:
     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()
+    backend_content = _backend_tokenizer_content(getattr(tokenizer, "backend_tokenizer", None))
+    if backend_content is not None:
+        content["backend"] = backend_content
+    return json_sha256(_fingerprint_jsonable(content), ensure_ascii=False)
 
 
 def _tokenizer_metadata(model: Any, tokenizer: Any | None) -> dict[str, Any]:
@@ -301,14 +322,12 @@ def _model_state_sha256(model: Any) -> str:
                 f"Cannot fingerprint meta-device model state entry {name!r}; pass "
                 "model_state_fingerprint with a caller-owned state identity."
             )
-        header = json.dumps(
+        header = compact_json(
             {
                 "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)
@@ -354,6 +373,7 @@ def _run_fingerprint(
     batch_size: int,
     batch_window_size: int,
     max_tokens_per_batch: int | None,
+    taps: Sequence[Mapping[str, Any]] | None = None,
 ) -> tuple[str, str, str | None, str]:
     input_fingerprint = _input_sha256(records)
     attention_backend = _attention_backend(model)
@@ -391,6 +411,7 @@ def _run_fingerprint(
         "execution": _execution_identity_metadata(model),
         "embedding_context": _fingerprint_jsonable(embedding_context),
         "pooling": list(pooling),
+        "pooling_semantics": dict(POOLING_SEMANTICS),
         "full_embeddings": full_embeddings,
         "max_length": max_length,
         "truncate": truncate,
@@ -399,16 +420,16 @@ def _run_fingerprint(
             "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()
+    if taps is not None:
+        # Each hidden tap binds its requested dtype; each reduced tap binds its own contract.
+        payload["taps"] = _fingerprint_jsonable(taps)
+    run_fingerprint = json_sha256(payload)
     return (
         input_fingerprint,
         run_fingerprint,
@@ -437,6 +458,7 @@ def _embedding_context(
     decoder_attention_mask: Tensor | None,
     model_kwargs: Mapping[str, Any],
 ) -> tuple[dict[str, Any], tuple[str, ...] | None]:
+    # decoder_input_ids, decoder_attention_mask: (n_records, l_decoder), aligned with records
     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)

fastplms/embeddings/inputs.py

@@ -3,46 +3,104 @@
 from __future__ import annotations
 
 import hashlib
+import shutil
 import sqlite3
 import tempfile
+import weakref
 
 from collections.abc import Iterable, Iterator, Mapping, Sequence
+from dataclasses import dataclass
 from pathlib import Path
-from typing import overload
+from typing import NamedTuple, 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."""
+class FastaRecord(NamedTuple):
+    """One FASTA record: its header text and its sequence lines joined."""
+
+    header: str
+    sequence: str
+
+
+@dataclass(frozen=True)
+class FastaDialect:
+    """How one caller reads FASTA lines. FastPLMs has three readers, and each keeps its rules here.
+
+    ``strip_lines``: strip each line before use; otherwise lines are used as given, so a space-only line
+        is sequence data.
+    ``comment_prefix``: lines that start with it are skipped, or ``None`` for no comments.
+    ``first_word_header``: the header is its first whitespace-delimited word, not the whole header text.
+    ``squeeze_sequence_whitespace``: remove all whitespace inside a sequence line.
+    ``orphan_message``: raised as ``ValueError`` for sequence data before the first header, formatted
+        with ``line_number`` and ``source``; ``None`` skips such lines.
+    ``empty_message``: raised as ``ValueError`` when the input holds no record, formatted with
+        ``source``; ``None`` allows empty input.
+    """
+
+    strip_lines: bool
+    comment_prefix: str | None
+    first_word_header: bool
+    squeeze_sequence_whitespace: bool
+    orphan_message: str | None
+    empty_message: str | None
+
 
-    identifier: str | None = None
+EMBEDDING_FASTA = FastaDialect(
+    strip_lines=True,
+    comment_prefix=None,
+    first_word_header=True,
+    squeeze_sequence_whitespace=True,
+    orphan_message="Sequence data precedes the first FASTA header on line {line_number}.",
+    empty_message="No FASTA records found in {source}.",
+)
+
+
+def scan_fasta_lines(
+    lines: Iterable[str], dialect: FastaDialect, *, source: str
+) -> Iterator[FastaRecord]:
+    """Yield the records of FASTA ``lines`` in order, one record at a time.
+
+    ``source`` names the input in the dialect's messages. A record is yielded when the next header or
+    the end of input shows it is complete, so an error raised later in the input surfaces after the
+    records before it.
+    """
+
+    header: 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:
+    for line_number, raw_line in enumerate(lines, start=1):
+        line = raw_line.strip() if dialect.strip_lines else raw_line
+        if not line or (dialect.comment_prefix is not None and line.startswith(dialect.comment_prefix)):
+            continue
+
+        if line.startswith(">"):
+            if header is not None:
+                found_record = True
+                yield FastaRecord(header, "".join(sequence_parts))
+            header = line[1:].strip()
+            if dialect.first_word_header:
+                # An empty header has no first word and raises IndexError here.
+                header = header.split(maxsplit=1)[0]
+            sequence_parts = []
+        elif header is not None:
+            sequence_parts.append("".join(line.split()) if dialect.squeeze_sequence_whitespace else line)
+        elif dialect.orphan_message is not None:
+            raise ValueError(dialect.orphan_message.format(line_number=line_number, source=source))
+
+    if header 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}.")
+        yield FastaRecord(header, "".join(sequence_parts))
+    if not found_record and dialect.empty_message is not None:
+        raise ValueError(dialect.empty_message.format(source=source))
+
+
+def iter_fasta(path: str | Path) -> Iterator[EmbeddingInput]:
+    """Yield FASTA records in source order without reading the file into memory."""
+
+    with Path(path).open("r", encoding="utf-8") as handle:
+        for record in scan_fasta_lines(handle, EMBEDDING_FASTA, source=str(path)):
+            yield EmbeddingInput(record.header, record.sequence)
 
 
 def parse_fasta(path: str | Path) -> list[EmbeddingInput]:
@@ -66,6 +124,27 @@ def _normalize_input_item(
     )
 
 
+class _SpoolFiles:
+    """A spool's directory and SQLite connection, released together: the connection first.
+
+    Windows cannot delete a file that an open connection holds. The cycle collector runs weakref
+    finalizers before any ``__del__``, so a spool freed as cyclic garbage, as one referenced from
+    a raised exception's traceback is, would have ``tempfile.TemporaryDirectory``'s own finalizer
+    remove the directory while the connection was still open. One finalizer owns both instead.
+    """
+
+    def __init__(self) -> None:
+        self.directory = Path(tempfile.mkdtemp(prefix="fastplms-inputs-"))
+        self.connection: sqlite3.Connection | None = None
+
+    def release(self) -> None:
+        if self.connection is not None:
+            self.connection.close()
+            self.connection = None
+        if self.directory.exists():
+            shutil.rmtree(self.directory)
+
+
 class _InputSpool(Sequence[EmbeddingInput]):
     """Immutable disk-backed normalized inputs with an incremental digest."""
 
@@ -73,19 +152,19 @@ class _InputSpool(Sequence[EmbeddingInput]):
         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)"
-        )
+        self._files = _SpoolFiles()
+        # Called by close(), or by garbage collection however the spool is freed; at most once.
+        self._release = weakref.finalize(self, self._files.release)
+        self.path = self._files.directory / "inputs.sqlite"
         digest = hashlib.sha256()
         count = 0
         pending: list[tuple[int, str, str]] = []
         try:
+            connection = self._files.connection = sqlite3.connect(self.path)
+            connection.execute(
+                "CREATE TABLE inputs ("
+                "position INTEGER PRIMARY KEY, input_id TEXT NOT NULL, sequence TEXT NOT NULL)"
+            )
             for position, item in enumerate(values):
                 record = _normalize_input_item(position, item)
                 for value in (record.id, record.sequence):
@@ -95,15 +174,15 @@ class _InputSpool(Sequence[EmbeddingInput]):
                 pending.append((position, record.id, record.sequence))
                 count += 1
                 if len(pending) == 1_024:
-                    self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
+                    connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
                     pending.clear()
             if pending:
-                self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
+                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(
+            connection.commit()
+            connection.close()
+            self._files.connection = sqlite3.connect(
                 f"{self.path.resolve().as_uri()}?mode=ro",
                 uri=True,
             )
@@ -115,9 +194,10 @@ class _InputSpool(Sequence[EmbeddingInput]):
         self._count = count
 
     def _require_connection(self) -> sqlite3.Connection:
-        if self._connection is None:
+        connection = self._files.connection
+        if connection is None:
             raise RuntimeError("Input spool is closed.")
-        return self._connection
+        return connection
 
     def __len__(self) -> int:
         return self._count
@@ -160,17 +240,7 @@ class _InputSpool(Sequence[EmbeddingInput]):
         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()
+        self._release()
 
 
 def _normalize_inputs(

fastplms/embeddings/output.py

@@ -125,11 +125,12 @@ class EmbeddingOutput:
             }
             self.sqlite_run_id = run_fingerprint
             if not resume and output_already_exists:
+                prior_run_readable = True
                 try:
                     load_sqlite_result(output, run_id=run_fingerprint)
                 except KeyError:
-                    pass
-                else:
+                    prior_run_readable = False
+                if prior_run_readable:
                     # Keep an exact prior run readable until replacement inference
                     # has produced the first complete commit window.
                     self.sqlite_replace_on_first_commit = True
@@ -209,7 +210,9 @@ class EmbeddingOutput:
             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)
+        embedding_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
+            return save_result(
+                embedding_result, self.output, format=self.format, shard_size=self.shard_size
+            )
+        return embedding_result

fastplms/embeddings/pooling.py

@@ -10,6 +10,63 @@ from torch import Tensor
 
 
 POOLING_NAMES = frozenset({"mean", "max", "norm", "median", "std", "var", "cls", "parti"})
+POOLING_SEMANTICS = {
+    "version": 2,
+    "mask": "biological_residues_only",
+    "accumulator": "float64_for_float64_input_else_float32",
+    "output_dtype": "input_dtype_after_reduction",
+    "variance_correction": 0,
+    "empty_residues": "reject",
+    "singleton_variance": 0,
+}
+
+
+# Pooling over every attended token, CLS and EOS included: the canonical policy of a store that
+# keeps the special tokens. Version 3 supersedes the residue-only version 2 for those stores.
+POOLING_SEMANTICS_TOKENS = {
+    "version": 3,
+    "mask": "all_attended_tokens_including_cls_and_eos",
+    "accumulator": "float64_for_float64_input_else_float32",
+    "output_dtype": "input_dtype_after_reduction",
+    "variance_correction": 0,
+    "empty_residues": "reject",
+    "singleton_variance": 0,
+}
+TOKEN_POOLING_NAMES = ("mean", "var", "std", "max", "norm")
+
+
+def pool_token_rows(X: Tensor, token_mask: Tensor, names: Sequence[str]) -> Tensor:
+    """Pool every attended token of each sequence, CLS and EOS included, without any device read.
+
+    The reductions equal ``Pooler`` over a mask that is true on every token: float32 accumulation
+    (float64 for float64 input), population variance from the two-pass mean, and a cast to the input
+    dtype after reduction. Unlike ``Pooler`` this skips its input validation, which reads device
+    values and stalls the host; the caller checks finiteness on the device after the copy lands.
+    """
+    # X: (b, n, d) padded, n = the longest l + 2; token_mask: (b, n) true on CLS, residues and EOS
+    unsupported = [name for name in names if name not in TOKEN_POOLING_NAMES]
+    if unsupported or not names or len(set(names)) != len(names):
+        raise ValueError(f"Token pooling supports {TOKEN_POOLING_NAMES}, each once; received {list(names)}.")
+    work = torch.float64 if X.dtype == torch.float64 else torch.float32
+    values = X.to(work)  # (b, n, d)
+    kept = token_mask.unsqueeze(-1)  # (b, n, 1)
+    count = kept.sum(dim=1).clamp_min(1).to(work)  # (b, 1), attended rows per sequence: l + 2
+    masked = values.masked_fill(~kept, 0)  # (b, n, d), padding zeroed
+    mean = masked.sum(dim=1) / count  # (b, d)
+    outputs: list[Tensor] = []
+    for name in names:
+        if name == "mean":
+            pooled = mean  # (b, d)
+        elif name == "max":
+            pooled = values.masked_fill(~kept, -torch.inf).amax(dim=1)  # (b, d)
+        elif name == "norm":
+            pooled = torch.linalg.vector_norm(masked, ord=2, dim=1)  # (b, d)
+        else:
+            centered = (values - mean.unsqueeze(1)).masked_fill(~kept, 0)  # (b, n, d)
+            variance = (centered * centered).sum(dim=1) / count  # (b, d)
+            pooled = variance.sqrt() if name == "std" else variance  # (b, d)
+        outputs.append(pooled.to(X.dtype))
+    return torch.cat(outputs, dim=-1)  # (b, len(names) * d)
 
 
 def _validate_inputs(X: Tensor, M: Tensor) -> Tensor:
@@ -42,6 +99,7 @@ def _pooled_attention(attentions: Tensor | Sequence[Tensor], *, batch_size: int)
     must not change that reduction.
     """
 
+    # attentions: (b, ..., l, l), or a sequence of (b, h, l, l) layer maps
     if isinstance(attentions, Sequence):
         if not attentions:
             raise ValueError("parti received an empty attention sequence.")
@@ -164,8 +222,12 @@ class Pooler:
         attentions: Tensor | Sequence[Tensor] | None = None,
         attention_backend: str | None = None,
     ) -> Tensor:
-        # X: (b, l, d); residue_mask: (b, l)
+        # X: (b, l, d); residue_mask: (b, l); attentions: (b, ..., l, l), or a sequence of (b, h, l, l) layer maps
         M = _validate_inputs(X, residue_mask)  # (b, l)
+        output_dtype = X.dtype
+        # Sum/variance in FP16 can overflow even when the final answer is representable.
+        # Retain FP64 precision, otherwise accumulate in FP32 and cast only the result.
+        X = X.to(dtype=torch.float64 if X.dtype == torch.float64 else torch.float32)  # (b, l, d)
         M_expanded = M.unsqueeze(-1)  # (b, l, 1)
         count = M_expanded.sum(dim=1).clamp_min(1)  # (b, 1)
         X_residues = X.masked_fill(~M_expanded, 0)  # (b, l, d)
@@ -206,6 +268,7 @@ class Pooler:
                     w = pagerank_weights(A_residue).to(dtype=X.dtype)  # (r,)
                     pooled.append(w @ X_i.index_select(0, indices))  # (d,)
                 Y = torch.stack(pooled)  # (b, d)
+            Y = Y.to(dtype=output_dtype)  # (b, d)
             if not bool(torch.isfinite(Y).all()):
                 raise ValueError(
                     f"Pooling operation {name!r} produced non-finite output from "
@@ -216,4 +279,7 @@ class Pooler:
         return torch.cat(outputs, dim=-1)  # (b, len(self.names) * d)
 
 
-__all__ = ["POOLING_NAMES", "Pooler", "pagerank_weights"]
+__all__ = [
+    "POOLING_NAMES", "POOLING_SEMANTICS_TOKENS", "TOKEN_POOLING_NAMES", "Pooler", "pagerank_weights",
+    "pool_token_rows",
+]

fastplms/embeddings/runner.py

@@ -12,6 +12,7 @@ from torch import Tensor
 from . import identity
 from .batches import (
     BatchExecutor,
+    TapExecutor,
     _residue_embeddings as _residue_embeddings,
     _temporary_eval,
     select_hidden_state_embeddings as select_hidden_state_embeddings,
@@ -36,8 +37,11 @@ from .inputs import (
     parse_fasta as parse_fasta,
 )
 from .output import EmbeddingOutput
-from .pooling import Pooler
-from .types import EmbeddingBatch, EmbeddingInput, EmbeddingResult
+from .pooling import POOLING_SEMANTICS, Pooler
+from .taps import Tap, TapPlan, plan_taps
+from .types import (
+    EmbeddingBatch, EmbeddingInput, EmbeddingResult, TapRecord, TapResult, TapRunReceipt,
+)
 
 
 _DEFAULT_BATCH_WINDOW_MULTIPLIER = 16
@@ -51,6 +55,9 @@ def embed_dataset(
     batch_size: int = 2,
     pooling: str | Sequence[str] | None = None,
     full_embeddings: bool = False,
+    taps: Sequence[Tap] | None = None,
+    tap_sink: Callable[[Sequence[TapRecord], Mapping[str, str]], None] | None = None,
+    require_residue_identity: bool = False,
     output: str | Path | None = None,
     format: str = "safetensors",
     resume: bool = True,
@@ -70,9 +77,16 @@ def embed_dataset(
     _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."""
+) -> EmbeddingResult | TapResult | TapRunReceipt:
+    """Embed protein sequences with stable ordering and residue-only pooling.
+
+    ``taps`` instead returns a ``TapResult``: each tap's output from one forward pass per
+    batch, kept in memory. With ``tap_sink``, deliver bounded windows to the callback and return
+    a ``TapRunReceipt`` without retaining their tensors. The callback receives ordered records
+    and the run/input fingerprints; it must release the records to preserve bounded memory.
+    """
 
+    # decoder_input_ids, decoder_attention_mask: (n_records, l_decoder), aligned with the inputs
     for name, value in (
         ("batch_size", batch_size),
         ("shard_size", shard_size),
@@ -94,11 +108,16 @@ def embed_dataset(
             raise ValueError(f"{optional_name} must be a positive integer when provided.")
     for name, value in (
         ("full_embeddings", full_embeddings),
+        ("require_residue_identity", require_residue_identity),
         ("resume", resume),
         ("truncate", truncate),
     ):
         if not isinstance(value, bool):
             raise TypeError(f"{name} must be a boolean.")
+    if require_residue_identity and taps is None:
+        raise ValueError("Residue identity validation requires a tap plan.")
+    if tap_sink is not None and (taps is None or not callable(tap_sink)):
+        raise ValueError("tap_sink requires a tap plan and a callable destination.")
     if not isinstance(format, str):
         raise TypeError("format must be a string.")
     if output is not None and not isinstance(output, (str, Path)):
@@ -135,15 +154,26 @@ def embed_dataset(
             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))
+    tap_plan = (
+        _tap_plan(
+            model,
+            taps,
+            pooling=pooling,
+            full_embeddings=full_embeddings,
+            output=output,
+            model_kwargs=model_kwargs,
+            family_adapter=(
+                _embedding_batch_fn is not None or _embedding_batch_identity is not None
+            ),
+        )
+        if taps is not None
+        else None
+    )
+    pooling_names = _requested_pooling(
+        pooling,
+        full_embeddings=full_embeddings,
+        taps_requested=tap_plan is not None,
     )
-    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:
@@ -187,22 +217,12 @@ def embed_dataset(
     )
     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)
+    records = _normalize_inputs(inputs, disk_backed=output is not None or tap_sink 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.")
@@ -215,7 +235,10 @@ def embed_dataset(
             f"by the model; unknown overrides: {sorted(unknown_pooling_overrides)}."
         )
     unsupported.difference_update(allowed_unsupported_pooling)
-    requested_unsupported = unsupported.intersection(pooling_names)
+    requested_pooling = set(pooling_names)
+    if tap_plan is not None:
+        requested_pooling.update(tap_plan.pooling_names)
+    requested_unsupported = unsupported.intersection(requested_pooling)
     if requested_unsupported:
         raise ValueError(
             f"{model.__class__.__name__} does not support pooling operations "
@@ -248,6 +271,7 @@ def embed_dataset(
         model.resolve_attn_implementation()
 
     tokenizer_metadata = _tokenizer_metadata(model, tokenizer)
+    tap_identity = tap_plan.identity() if tap_plan is not None else None
     (
         input_fingerprint,
         run_fingerprint,
@@ -269,7 +293,64 @@ def embed_dataset(
         batch_size=batch_size,
         batch_window_size=resolved_batch_window_size,
         max_tokens_per_batch=max_tokens_per_batch,
+        taps=tap_identity,
     )
+    attention_backend = _attention_backend(model)
+    if tap_plan is not None:
+        tap_records, tap_pool_slices = _embed_tap_windows(
+            model,
+            records,
+            TapExecutor(
+                model=model,
+                plan=tap_plan,
+                batch_size=batch_size,
+                max_tokens_per_batch=max_tokens_per_batch,
+                max_length=max_length,
+                truncate=truncate,
+                tokenizer=tokenizer,
+                dtype=dtype,
+                attention_backend=attention_backend,
+                require_residue_identity=require_residue_identity,
+            ),
+            window_size=resolved_batch_window_size,
+            sink=tap_sink,
+            run_identity={
+                "run_fingerprint": run_fingerprint, "input_fingerprint": input_fingerprint,
+            },
+        )
+        metadata = _run_metadata(
+            model,
+            records,
+            run_fingerprint=run_fingerprint,
+            input_fingerprint=input_fingerprint,
+            model_state_fingerprint=resolved_model_state_fingerprint,
+            model_state_fingerprint_source=model_state_fingerprint_source,
+            dtype=dtype,
+            attention_backend=attention_backend,
+            layer=None,
+            tokenizer_metadata=tokenizer_metadata,
+            embedding_context=embedding_context,
+            pooling_names=pooling_names,
+            pool_slices={},
+            full_embeddings=full_embeddings,
+            max_length=max_length,
+            truncate=truncate,
+            batch_size=batch_size,
+            batch_window_size=resolved_batch_window_size,
+            max_tokens_per_batch=max_tokens_per_batch,
+            output=output,
+            format=format,
+            descriptor_index="not-recorded",
+            taps={
+                "plan": tap_identity,
+                "stop_after_layer": tap_plan.deepest_layer,
+                "pool_slices": tap_pool_slices,
+            },
+        )
+        if tap_sink is not None:
+            metadata["storage_format"] = "tap-sink"
+            return TapRunReceipt(len(records), metadata)
+        return TapResult(tap_records, metadata)
     destination = EmbeddingOutput(
         records,
         output=output,
@@ -286,7 +367,6 @@ def embed_dataset(
     if destination.completed is not None:
         return destination.completed
 
-    attention_backend = _attention_backend(model)
     executor = BatchExecutor(
         model=model,
         batch_size=batch_size,
@@ -321,13 +401,181 @@ def embed_dataset(
             )
             destination.append(window_start, new_records)
 
-    software_versions = identity._software_versions()
-    projection = getattr(model, "embedding_projection", None)
-    resolved_layer = getattr(
+    metadata = _run_metadata(
         model,
-        "embedding_layer",
-        model_kwargs.get("hidden_state_index", -1),
+        records,
+        run_fingerprint=run_fingerprint,
+        input_fingerprint=input_fingerprint,
+        model_state_fingerprint=resolved_model_state_fingerprint,
+        model_state_fingerprint_source=model_state_fingerprint_source,
+        dtype=dtype,
+        attention_backend=attention_backend,
+        layer=getattr(
+            model,
+            "embedding_layer",
+            model_kwargs.get("hidden_state_index", -1),
+        ),
+        tokenizer_metadata=tokenizer_metadata,
+        embedding_context=embedding_context,
+        pooling_names=pooling_names,
+        pool_slices=pool_slices,
+        full_embeddings=full_embeddings,
+        max_length=max_length,
+        truncate=truncate,
+        batch_size=batch_size,
+        batch_window_size=resolved_batch_window_size,
+        max_tokens_per_batch=max_tokens_per_batch,
+        output=output,
+        format=format,
+        descriptor_index=(
+            "memory-metadata"
+            if output is None
+            else "sqlite-records"
+            if format == "sqlite"
+            else "safetensors-generation-index"
+        ),
     )
+    if destination.output_descriptors is not None:
+        metadata["outputs"] = destination.output_descriptors
+        metadata["tensor_hashes"] = [item["sha256"] for item in destination.output_descriptors]
+    return destination.finish(metadata)
+
+
+def _tap_plan(
+    model: Any,
+    taps: Sequence[Tap],
+    *,
+    pooling: str | Sequence[str] | None,
+    full_embeddings: bool,
+    output: str | Path | None,
+    model_kwargs: Mapping[str, Any],
+    family_adapter: bool,
+) -> TapPlan:
+    """Validate a tap request against the arguments it excludes and the model's hidden states."""
+
+    excluded = [
+        name
+        for name, requested in (
+            ("pooling", pooling is not None),
+            ("full_embeddings", full_embeddings),
+            ("hidden_state_index", "hidden_state_index" in model_kwargs),
+            ("store_all_hidden_states", "store_all_hidden_states" in model_kwargs),
+        )
+        if requested
+    ]
+    if excluded:
+        raise ValueError(
+            f"taps= cannot be combined with {', '.join(excluded)}; each tap names its own "
+            "layer and pooling."
+        )
+    if model_kwargs:
+        raise ValueError(
+            f"taps= takes no model keyword arguments; received {sorted(model_kwargs)}."
+        )
+    if family_adapter:
+        raise ValueError(
+            "taps= runs the model's own one-pass path; _embedding_batch_fn and "
+            "_embedding_batch_identity do not apply."
+        )
+    if output is not None:
+        raise ValueError(
+            "taps= returns its records in memory; omit output=. To persist them, call "
+            "embed_into_features, which writes each tap into the feature store of its key and "
+            "embeds only the sequences that store lacks."
+        )
+    if getattr(model, "embedding_tap_support", False) is not True:
+        raise ValueError(
+            f"{model.__class__.__name__} does not support taps=. One-pass taps need a model "
+            "family that implements them, such as ESM++ (ESMC)."
+        )
+    return plan_taps(taps, int(model.embedding_tap_state_count))
+
+
+def _requested_pooling(
+    pooling: str | Sequence[str] | None,
+    *,
+    full_embeddings: bool,
+    taps_requested: bool,
+) -> tuple[str, ...]:
+    """Pooler names of a single-output run: mean by default, none for residues or taps."""
+
+    if taps_requested:
+        return ()
+    if full_embeddings:
+        if pooling is not None:
+            raise ValueError("full_embeddings=True cannot be combined with pooling.")
+        return ()
+    names = (
+        ("mean",)
+        if pooling is None
+        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
+    )
+    if not names:
+        raise ValueError("pooling is required unless full_embeddings=True.")
+    return names
+
+
+def _embed_tap_windows(
+    model: Any,
+    records: Sequence[EmbeddingInput],
+    executor: TapExecutor,
+    *,
+    window_size: int,
+    sink: Callable[[Sequence[TapRecord], Mapping[str, str]], None] | None = None,
+    run_identity: Mapping[str, str] | None = None,
+) -> tuple[list[TapRecord], dict[str, dict[str, tuple[int, int]]]]:
+    """Run bounded windows in source order, retaining tensors only without a sink."""
+
+    tap_records: list[TapRecord] = []
+    pool_slices: dict[str, dict[str, tuple[int, int]]] = {}
+    with _temporary_eval(model), torch.inference_mode():
+        for window_start in range(0, len(records), window_size):
+            window_stop = min(window_start + 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
+            )
+            if sink is None:
+                tap_records.extend(new_records)
+            else:
+                sink(new_records, dict(run_identity or {}))
+            # Release the previous window before allocating the next, including on CPU.
+            del new_records
+    return tap_records, pool_slices
+
+
+def _run_metadata(
+    model: Any,
+    records: Sequence[EmbeddingInput],
+    *,
+    run_fingerprint: str,
+    input_fingerprint: str,
+    model_state_fingerprint: str | None,
+    model_state_fingerprint_source: str,
+    dtype: torch.dtype | None,
+    attention_backend: str | None,
+    layer: Any,
+    tokenizer_metadata: dict[str, Any],
+    embedding_context: Mapping[str, Any],
+    pooling_names: Sequence[str],
+    pool_slices: Mapping[str, tuple[int, int]],
+    full_embeddings: bool,
+    max_length: int | None,
+    truncate: bool,
+    batch_size: int,
+    batch_window_size: int,
+    max_tokens_per_batch: int | None,
+    output: str | Path | None,
+    format: str,
+    descriptor_index: str,
+    taps: Mapping[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Everything a finished run records so that it can be reproduced and resumed."""
+
+    software_versions = identity._software_versions()
+    projection = getattr(model, "embedding_projection", None)
     token_policy = getattr(
         model,
         "embedding_token_policy",
@@ -349,14 +597,14 @@ def embed_dataset(
         "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": 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,
+        "layer": layer,
         "projection": projection,
         "esmc_source": getattr(model, "_esmc_source", None),
         "esmc_revision": getattr(model, "_esmc_source_revision", None),
@@ -365,14 +613,18 @@ def embed_dataset(
         "tokenizer": tokenizer_metadata,
         **embedding_context,
         "pooling": list(pooling_names),
+        "pooling_semantics": dict(POOLING_SEMANTICS),
         "pool_slices": pool_slices,
         "full_embeddings": full_embeddings,
         "max_length": max_length,
         "truncate": truncate,
         "truncation": {"enabled": truncate, "max_length": max_length},
+        "retained_positions": (
+            "biological_residues_in_input_order_after_optional_prefix_crop_before_forward"
+        ),
         "batching": {
             "batch_size": batch_size,
-            "batch_window_size": resolved_batch_window_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"),
             "ordering": "bounded-length-bucketed-stable-output",
@@ -386,13 +638,7 @@ def embed_dataset(
         },
         "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"
-        ),
+        "descriptor_index": descriptor_index,
         "storage_format": format if output is not None else "memory",
         "software": software_versions,
         "execution": _execution_identity_metadata(model),
@@ -401,19 +647,20 @@ def embed_dataset(
         "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]
+    if taps is not None:
+        metadata["taps"] = taps
     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)
+    return metadata
 
 
 class EmbeddingMixin:
     """Small delegation mixin shared by FastPLMs model classes."""
 
-    def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult:
+    def embed_dataset(
+        self, inputs: Any, **kwargs: Any,
+    ) -> EmbeddingResult | TapResult | TapRunReceipt:
         return embed_dataset(self, inputs, **kwargs)
 
 

fastplms/embeddings/storage.py

@@ -22,6 +22,7 @@ from .types import (
     EmbeddingResult,
     LazyTensorReference,
 )
+from ..json_files import indented_json
 
 
 _DTYPE_NAMES: dict[torch.dtype, str] = {
@@ -102,13 +103,15 @@ def _bounded_tensor_chunks(X: Tensor, max_bytes: int) -> Iterator[Tensor]:
 
 
 def _tensor_hash_chunks(X: Tensor) -> Iterator[bytes]:
-    for chunk in _bounded_tensor_chunks(X, _TENSOR_HASH_CHUNK_BYTES):
+    # X: (...)
+    for chunk in _bounded_tensor_chunks(X, _TENSOR_HASH_CHUNK_BYTES):  # (n_chunk,)
         yield chunk.view(torch.uint8).numpy().tobytes()
 
 
 def tensor_sha256(X: Tensor) -> str:
     """Hash dtype, shape, and exact tensor bytes."""
 
+    # X: (...)
     if not isinstance(X, Tensor):
         raise TypeError("X must be a tensor.")
     if X.dtype not in _DTYPE_NAMES:
@@ -126,22 +129,23 @@ def tensor_sha256(X: Tensor) -> str:
 
 
 def _encode_tensor(X: Tensor) -> tuple[str, str, bytes]:
+    # X: (...)
     if X.dtype not in _DTYPE_NAMES:
         raise TypeError(f"Unsupported tensor dtype {X.dtype}.")
     shape = json.dumps(tuple(X.shape), separators=(",", ":"))
     return _DTYPE_NAMES[X.dtype], shape, _tensor_bytes(X)
 
 
-def _decode_tensor(dtype_name: str, shape_json: str, data: bytes) -> Tensor:
+def _decode_tensor(dtype_name: str, shape_json: str, raw_bytes: bytes) -> Tensor:
     try:
         dtype = _NAME_DTYPES[dtype_name]
     except KeyError as error:
         raise ValueError(f"Unsupported stored dtype {dtype_name!r}.") from error
     shape = tuple(json.loads(shape_json))
     # uint8 is used only as a byte-level carrier, preserving BF16 bits exactly.
-    byte_array = np.frombuffer(data, dtype=np.uint8).copy()  # (n_bytes,)
+    byte_array = np.frombuffer(raw_bytes, dtype=np.uint8).copy()  # (n_bytes,)
     X = torch.from_numpy(byte_array).view(dtype)  # (n_elements,)
-    return X.reshape(shape).clone()  # shape
+    return X.reshape(shape).clone()  # (...), the stored shape
 
 
 def _index_path(path: str | Path) -> Path:
@@ -172,10 +176,6 @@ def _resolve_index_child(root: Path, relative: str, *, label: str) -> Path:
     return candidate
 
 
-def _canonical_json_bytes(payload: dict[str, Any]) -> bytes:
-    return (json.dumps(payload, indent=2, sort_keys=True) + "\n").encode("utf-8")
-
-
 def _load_authoritative_index(
     path: str | Path,
 ) -> tuple[dict[str, Any], Path, dict[str, Any]]:
@@ -198,7 +198,7 @@ def _load_authoritative_index(
         snapshot = run_manifest.get("index_payload")
         if isinstance(snapshot, dict):
             payload = snapshot
-            index_bytes = _canonical_json_bytes(payload)
+            index_bytes = indented_json(payload).encode("utf-8")
         elif snapshot is None:
             index_bytes = stable_index_path.read_bytes()
             payload = json.loads(index_bytes.decode("utf-8"))
@@ -268,7 +268,7 @@ def _load_safetensor(path: Path, key: str) -> Tensor:
     except ImportError as error:
         raise ImportError("Loading embeddings requires the 'safetensors' package.") from error
     with safe_open(path, framework="pt", device="cpu") as handle:
-        return cast(Tensor, handle.get_tensor(key))
+        return cast(Tensor, handle.get_tensor(key))  # (...), as stored under key
 
 
 def _safetensors_shard_prefix(path: str | Path) -> str:
@@ -361,7 +361,7 @@ def _record_from_safetensors_descriptor(root: Path, item: dict[str, Any]) -> Emb
         raise ValueError(f"Safetensors tensor shard is missing: {relative}.")
 
     def load_tensor() -> Tensor:
-        return _load_safetensor(tensor_path, key)
+        return _load_safetensor(tensor_path, key)  # (...), the descriptor's shape
 
     reference = LazyTensorReference(
         source=str(tensor_path),
@@ -720,7 +720,7 @@ class SafetensorsStreamWriter:
             raise FileExistsError(
                 f"Refusing to reuse immutable safetensors generation index {generation_index_path}."
             )
-        encoded_index = _canonical_json_bytes(payload)
+        encoded_index = indented_json(payload).encode("utf-8")
         temporary_generation_index.write_bytes(encoded_index)
         temporary_generation_index.replace(generation_index_path)
 
@@ -739,7 +739,7 @@ class SafetensorsStreamWriter:
         temporary_manifest = self.run_manifest_path.with_name(
             f".{self.run_manifest_path.name}.{pointer_identity}.tmp"
         )
-        temporary_manifest.write_bytes(_canonical_json_bytes(run_manifest))
+        temporary_manifest.write_bytes(indented_json(run_manifest).encode("utf-8"))
         temporary_manifest.replace(self.run_manifest_path)
 
         # ``index.json`` is a non-authoritative convenience pointer. The run
@@ -753,7 +753,7 @@ class SafetensorsStreamWriter:
         temporary_index = self.index_path.with_name(
             f".{self.index_path.name}.{pointer_identity}.tmp"
         )
-        temporary_index.write_bytes(_canonical_json_bytes(stable_pointer))
+        temporary_index.write_bytes(indented_json(stable_pointer).encode("utf-8"))
         temporary_index.replace(self.index_path)
 
         return load_safetensors_result(self.index_path)
@@ -771,7 +771,7 @@ class SafetensorsStreamWriter:
 
 
 def save_safetensors_result(
-    result: EmbeddingResult,
+    embedding_result: EmbeddingResult,
     path: str | Path,
     *,
     shard_size: int = DEFAULT_SHARD_SIZE,
@@ -780,13 +780,13 @@ def save_safetensors_result(
 
     writer = SafetensorsStreamWriter(
         path,
-        result.metadata,
+        embedding_result.metadata,
         shard_size=shard_size,
         publish_initial=False,
         publish_incremental=False,
     )
-    writer.append(result, publish=False)
-    return writer.publish(complete=bool(result.metadata.get("complete", True)))
+    writer.append(embedding_result, publish=False)
+    return writer.publish(complete=bool(embedding_result.metadata.get("complete", True)))
 
 
 def load_safetensors_result(path: str | Path) -> EmbeddingResult:
@@ -925,19 +925,19 @@ def _ensure_sqlite_schema(connection: sqlite3.Connection) -> None:
         connection.commit()
 
 
-def save_sqlite_result(result: EmbeddingResult, path: str | Path) -> EmbeddingResult:
+def save_sqlite_result(embedding_result: EmbeddingResult, path: str | Path) -> EmbeddingResult:
     """Transactionally store an ordered result in normalized SQLite tables."""
 
     path = Path(path)
     path.parent.mkdir(parents=True, exist_ok=True)
-    run_id = str(result.metadata.get("run_fingerprint", ""))
+    run_id = str(embedding_result.metadata.get("run_fingerprint", ""))
     if not run_id:
         raise ValueError("SQLite results require metadata['run_fingerprint'].")
     metadata_json = json.dumps(
         _persistent_metadata(
-            result.metadata,
+            embedding_result.metadata,
             descriptor_index="sqlite-records",
-            record_count=len(result),
+            record_count=len(embedding_result),
         ),
         sort_keys=True,
     )
@@ -951,13 +951,13 @@ def save_sqlite_result(result: EmbeddingResult, path: str | Path) -> EmbeddingRe
             "SELECT ?, ?, COALESCE(MAX(published_order), 0) + 1 FROM runs",
             (run_id, metadata_json),
         )
-        for position, record in enumerate(result):
+        for position, record in enumerate(embedding_result):
             X = record.load_tensor().detach().cpu().contiguous()  # (...)
-            dtype_name, shape_json, data = _encode_tensor(X)
+            dtype_name, shape_json, raw_bytes = _encode_tensor(X)
             digest = tensor_sha256(X)
             connection.execute(
                 "INSERT INTO tensors VALUES (?, ?, ?, ?, ?, ?)",
-                (run_id, position, dtype_name, shape_json, data, digest),
+                (run_id, position, dtype_name, shape_json, raw_bytes, digest),
             )
             connection.execute(
                 "INSERT INTO records VALUES (?, ?, ?, ?)",
@@ -1057,11 +1057,11 @@ def append_sqlite_records(
         for offset, record in enumerate(records):
             position = start_position + offset
             X = record.load_tensor().detach().cpu().contiguous()  # (...)
-            dtype_name, shape_json, data = _encode_tensor(X)
+            dtype_name, shape_json, raw_bytes = _encode_tensor(X)
             digest = tensor_sha256(X)
             connection.execute(
                 "INSERT INTO tensors VALUES (?, ?, ?, ?, ?, ?)",
-                (run_id, position, dtype_name, shape_json, data, digest),
+                (run_id, position, dtype_name, shape_json, raw_bytes, digest),
             )
             connection.execute(
                 "INSERT INTO records VALUES (?, ?, ?, ?)",
@@ -1142,7 +1142,7 @@ def _load_sqlite_tensor(path: Path, run_id: str, position: int) -> Tensor:
         ).fetchone()
     if row is None:
         raise KeyError(f"Missing SQLite tensor {run_id}:{position}.")
-    return _decode_tensor(*row)
+    return _decode_tensor(*row)  # (...), the stored shape
 
 
 def _validate_sqlite_descriptor_row(
@@ -1180,7 +1180,7 @@ def _sqlite_record_from_row(path: Path, run_id: str, row: Sequence[Any]) -> Embe
     )
 
     def load_tensor() -> Tensor:
-        return _load_sqlite_tensor(path, run_id, position)
+        return _load_sqlite_tensor(path, run_id, position)  # (...), the stored shape
 
     reference = LazyTensorReference(
         source=str(path),
@@ -1284,7 +1284,7 @@ def load_sqlite_result(
         _validate_sqlite_result_schema(connection, path)
         if run_id is None:
             run_columns = {
-                str(info[1]) for info in connection.execute("PRAGMA table_info(runs)").fetchall()
+                str(column[1]) for column in connection.execute("PRAGMA table_info(runs)").fetchall()
             }
             if "published_order" in run_columns:
                 row = connection.execute(
@@ -1434,36 +1434,36 @@ _LEGACY_CODE_DTYPES: dict[int, tuple[np.dtype[Any], torch.dtype]] = {
 
 
 def _decode_legacy_sqlite_blob(
-    data: bytes,
+    blob: bytes,
     *,
     fallback_shape: tuple[int, ...] | None,
     allow_unsafe_pickle: bool,
 ) -> Tensor:
-    if len(data) >= 6 and data[0] == _LEGACY_COMPACT_VERSION:
-        dtype_code = int(data[1])
+    if len(blob) >= 6 and blob[0] == _LEGACY_COMPACT_VERSION:
+        dtype_code = int(blob[1])
         if dtype_code not in _LEGACY_CODE_DTYPES:
             raise ValueError(f"Unsupported legacy compact dtype code {dtype_code}.")
-        (ndim,) = struct.unpack_from("<i", data, 2)
-        if ndim < 0 or ndim > 16 or len(data) < 6 + 4 * ndim:
+        (ndim,) = struct.unpack_from("<i", blob, 2)
+        if ndim < 0 or ndim > 16 or len(blob) < 6 + 4 * ndim:
             raise ValueError("Malformed legacy compact embedding header.")
-        shape = tuple(int(value) for value in struct.unpack_from(f"<{ndim}i", data, 6))
+        shape = tuple(int(value) for value in struct.unpack_from(f"<{ndim}i", blob, 6))
         if any(size < 0 for size in shape):
             raise ValueError("Malformed negative legacy embedding dimension.")
         numpy_dtype, target_dtype = _LEGACY_CODE_DTYPES[dtype_code]
         offset = 6 + 4 * ndim
         expected = int(np.prod(shape, dtype=np.int64)) * numpy_dtype.itemsize
-        if len(data) - offset != expected:
+        if len(blob) - offset != expected:
             raise ValueError("Legacy compact embedding payload length does not match shape.")
-        array = (  # shape
-            np.frombuffer(data, dtype=numpy_dtype, offset=offset).copy().reshape(shape)
+        array = (  # (...), the stored shape
+            np.frombuffer(blob, dtype=numpy_dtype, offset=offset).copy().reshape(shape)
         )
-        return torch.from_numpy(array).to(dtype=target_dtype)  # shape
+        return torch.from_numpy(array).to(dtype=target_dtype)  # (...), the stored shape
 
     try:
-        loaded = torch.load(io.BytesIO(data), map_location="cpu", weights_only=True)
+        loaded = torch.load(io.BytesIO(blob), map_location="cpu", weights_only=True)
     except Exception as safe_error:
         if allow_unsafe_pickle:
-            loaded = torch.load(io.BytesIO(data), map_location="cpu", weights_only=False)
+            loaded = torch.load(io.BytesIO(blob), map_location="cpu", weights_only=False)
         elif fallback_shape is None:
             raise ValueError(
                 "Legacy embedding blob is neither compact nor safely loadable. "
@@ -1472,14 +1472,14 @@ def _decode_legacy_sqlite_blob(
             ) from safe_error
         else:
             expected = int(np.prod(fallback_shape, dtype=np.int64)) * 4
-            if len(data) != expected:
+            if len(blob) != expected:
                 raise ValueError(
                     "Legacy raw FP32 payload length does not match fallback_shape."
                 ) from safe_error
-            array = np.frombuffer(data, dtype=np.float32).copy().reshape(  # fallback_shape
+            array = np.frombuffer(blob, dtype=np.float32).copy().reshape(  # (...), fallback_shape
                 fallback_shape
             )
-            return torch.from_numpy(array)  # fallback_shape
+            return torch.from_numpy(array)  # (...), fallback_shape
     if not isinstance(loaded, Tensor):
         raise ValueError("Legacy serialized embedding payload must contain one tensor.")
     return loaded.detach().cpu()  # (...)
@@ -1522,13 +1522,13 @@ def convert_legacy_sqlite(
 
     records: list[EmbeddingRecord] = []
     content_digest = hashlib.sha256()
-    for position, (sequence, data) in enumerate(rows):
+    for position, (sequence, blob) in enumerate(rows):
         if not isinstance(sequence, str) or not sequence:
             raise ValueError("Legacy embedding sequences must be non-empty strings.")
-        if not isinstance(data, bytes):
-            data = bytes(data)
+        if not isinstance(blob, bytes):
+            blob = bytes(blob)
         tensor = _decode_legacy_sqlite_blob(
-            data,
+            blob,
             fallback_shape=fallback_shape,
             allow_unsafe_pickle=allow_unsafe_pickle,
         )
@@ -1559,16 +1559,16 @@ def convert_legacy_sqlite(
 
 
 def save_result(
-    result: EmbeddingResult,
+    embedding_result: EmbeddingResult,
     path: str | Path,
     *,
     format: str = "safetensors",
     shard_size: int = DEFAULT_SHARD_SIZE,
 ) -> EmbeddingResult:
     if format == "safetensors":
-        return save_safetensors_result(result, path, shard_size=shard_size)
+        return save_safetensors_result(embedding_result, path, shard_size=shard_size)
     if format == "sqlite":
-        return save_sqlite_result(result, path)
+        return save_sqlite_result(embedding_result, path)
     if format == "pth":
         raise ValueError("Writing pickle-based .pth embeddings is not supported.")
     raise ValueError("format must be 'safetensors' or 'sqlite'.")

fastplms/embeddings/taps.py

@@ -0,0 +1,413 @@
+"""Tap plans: several hidden-state outputs from one forward pass per batch.
+
+A tap names one hidden state and what to keep from it. ``HiddenTap`` keeps the rows the mask selects
+(biological residues in a residue run; CLS, residues and EOS in a token run) or pools them
+(``Pooler`` in a residue run, ``pool_token_rows`` in a token run). ``ReducedTap`` hands the state to a
+caller's reducer, such as a sparse-autoencoder encoder and its pooling. ``StreamingTap`` reduces layers
+as they arrive without saving their full hidden states. A plan runs one forward pass per
+batch that stops once the deepest tapped state exists.
+
+Layer indices follow the FastPLMs hidden-state order: index ``i`` is the input to block ``i``,
+index ``n`` (the block count) is the final normalized state, and negative indices count back
+from it, so ``-1`` is the final state.
+
+Symbols: b sequences of a batch; l token columns of the padded batch (CLS, residues, EOS, padding); d hidden
+width; n attended rows of a batch; r residues of one sequence. In a residue run the mask is false on CLS, EOS and
+padding, so a sequence is r rows of an ``(n, d)`` output. In a token run (canonical) it is false on padding only, so a
+sequence is r + 2 rows (row 0 CLS, rows 1..r residues, row r + 1 EOS) in every ``(n, d)`` output and every pooling.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import torch
+
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+from torch import Tensor
+
+from types import MappingProxyType
+from .pooling import Pooler
+from ..features.layouts import TopKRow
+
+
+@dataclass(frozen=True, slots=True)
+class RowSelection:
+    """The rows of a ``(b, l)`` grid that a tap keeps, known on the host so no device value is read.
+
+    ``flat_index`` holds each kept row's position in the row-major ``(b * l)`` grid, ``owner`` the
+    sequence it belongs to, and ``counts`` how many rows each sequence keeps. Gathering with an index
+    built from host-known lengths never stalls the host the way boolean indexing does.
+    """
+
+    flat_index: Tensor  # (n,) int64 on the device, n = sum of counts
+    owner: Tensor  # (n,) int64 on the device, the sequence of each kept row
+    counts: tuple[int, ...]  # (b,) kept rows per sequence, on the host
+    sizes: Tensor  # (b,) int64 on the device, the same counts
+
+    def gather(self, X: Tensor) -> Tensor:
+        """The kept rows of ``X`` in sequence order, without padding."""
+        # X: (b, l, w) -> (n, w)
+        return X.reshape(-1, X.shape[-1]).index_select(0, self.flat_index)  # (n, w)
+
+
+@dataclass(frozen=True, slots=True)
+class TapBatch:
+    """One batch of a tapped hidden state, as a ``ReducedTap`` reducer receives it.
+
+    ``X`` has shape ``(b, l, d)``. ``token_mask`` has shape ``(b, l)`` and marks every attended
+    token, BOS and EOS included. ``residue_mask`` has shape ``(b, l)`` and marks the rows that
+    ``HiddenTap`` keeps and pools: the biological residues, or, when the run keeps the special
+    tokens, every attended token (``l`` then counts CLS and EOS). All three share X's device.
+    ``rows`` is the host-known selection of ``residue_mask`` when the executor has one, and
+    ``cache`` lets taps of one batch share an intermediate such as an SAE encoding.
+    """
+
+    X: Tensor
+    token_mask: Tensor
+    residue_mask: Tensor
+    rows: RowSelection | None = None
+    cache: dict[Any, Any] = field(default_factory=dict)
+
+    def selection(self) -> RowSelection:
+        """The kept rows, taken from the executor when it supplied them, else read from the mask."""
+        if self.rows is not None:
+            return self.rows
+        kept = self.residue_mask.bool()  # (b, l)
+        owner, position = kept.nonzero(as_tuple=True)  # (n,), (n,)
+        sizes = kept.sum(dim=1)  # (b,)
+        return RowSelection(owner * kept.shape[1] + position, owner, tuple(sizes.tolist()), sizes)
+
+
+@dataclass(frozen=True, slots=True)
+class HiddenTap:
+    """One hidden state, kept as ragged rows or pooled.
+
+    ``pooling=None`` keeps one ``(r_i, d)`` tensor of mask-selected rows per sequence: ``r_i`` biological
+    residues in a residue run, ``l_i + 2`` rows (CLS, residues, EOS) in a token run. Pooler
+    names give one pooled vector per sequence over those same rows, concatenated in request order. ``parti`` needs the
+    attention graph of a full pass, so a tap rejects it. ``dtype`` overrides the run's dtype
+    for this output, starting from the original captured state; None inherits the run dtype.
+    """
+
+    name: str
+    layer: int
+    pooling: str | Sequence[str] | None = None
+    dtype: torch.dtype | None = None
+
+    def __post_init__(self) -> None:
+        _require_name_and_layer(self.name, self.layer)
+        if self.dtype is not None and self.dtype not in (
+            torch.float16, torch.bfloat16, torch.float32, torch.float64
+        ):
+            raise ValueError(
+                "A hidden tap dtype must be float16, bfloat16, float32, float64, "
+                "or None to inherit the run dtype."
+            )
+        if self.pooling is None:
+            return
+        names = Pooler(self.pooling).names  # validates names and rejects duplicates
+        if "parti" in names:
+            raise ValueError(
+                f"Tap {self.name!r} cannot pool with 'parti', which needs the attention graph "
+                "of a full forward pass."
+            )
+        object.__setattr__(self, "pooling", names)
+
+
+@dataclass(frozen=True, slots=True)
+class ReducedTap:
+    """One hidden state reduced by a caller's function.
+
+    ``reduce`` maps a ``TapBatch`` to a tensor with one row per sequence, shape ``(b, ...)``.
+    ``identity`` describes the reducer in plain data: strings, numbers, booleans, None, lists,
+    and string-keyed mappings. The run fingerprint records it, so two reducers with equal
+    identities must return equal outputs.
+    """
+
+    name: str
+    layer: int
+    reduce: Callable[[TapBatch], Tensor]
+    identity: Mapping[str, Any]
+
+    def __post_init__(self) -> None:
+        _require_name_and_layer(self.name, self.layer)
+        if not callable(self.reduce):
+            raise TypeError(f"Tap {self.name!r} needs a callable reduce.")
+        if not isinstance(self.identity, Mapping) or not self.identity:
+            raise TypeError(f"Tap {self.name!r} needs a non-empty identity mapping.")
+        _require_plain_data(self.identity, f"identity of tap {self.name!r}")
+        # A private canonical copy, so later changes to the caller's mapping cannot change the
+        # fingerprint of this tap.
+        canonical = json.loads(json.dumps(dict(self.identity), sort_keys=True, allow_nan=False))
+        object.__setattr__(self, "identity", MappingProxyType(canonical))
+
+
+@dataclass(frozen=True, slots=True)
+class SparseResidueTap:
+    """Reduce one state to sparse codes per kept row, in sequence and row order.
+
+    The reducer receives the same masks as dense taps and returns one ``TopKRow`` per sequence.
+    Its output must retain every row the mask keeps: the biological residues in a residue run, all
+    ``l_i + 2`` attended tokens in a token run. Both integer indices and floating
+    values remain sparse through extraction and persistence.
+    """
+
+    name: str
+    layer: int
+    reduce: Callable[[TapBatch], Sequence[TopKRow]]
+    identity: Mapping[str, Any]
+    codebook_size: int
+    sparse_count: int
+    # A run that keeps CLS and EOS reads every sequence's codes as one packed (n, k) pair, n = sum(l_i + 2),
+    # and splits nothing: the token executor calls this instead of ``reduce``. Not part of the identity.
+    reduce_packed: Callable[[TapBatch], TopKRow] | None = None
+
+    def __post_init__(self) -> None:
+        checked = ReducedTap(self.name, self.layer, self.reduce, self.identity)
+        object.__setattr__(self, "identity", checked.identity)
+        if (type(self.codebook_size) is not int or not 1 <= self.codebook_size <= 2**31
+                or type(self.sparse_count) is not int
+                or not 1 <= self.sparse_count <= self.codebook_size):
+            raise ValueError(
+                "Sparse residue taps require integer 1 <= sparse_count <= codebook_size <= 2**31."
+            )
+
+
+class LayerAccumulator(Protocol):
+    """Batch-local state for a streaming reduction. Never mutate the borrowed hidden state."""
+
+    def update(self, layer: int, batch: TapBatch) -> None: ...
+
+    def finish(self) -> Tensor:
+        """Return a token-aligned tensor of shape (b, l, c)."""
+        ...
+
+
+@dataclass(frozen=True, slots=True)
+class StreamingTap:
+    """Reduce selected layers as they arrive, retaining only the accumulator's own state.
+
+    ``begin`` creates a fresh accumulator for each batch. ``update`` borrows each original
+    hidden state, before the run's output dtype conversion, in ascending layer order. ``finish``
+    returns token-aligned residue features; the engine applies its biological mask and restores
+    input order. Reducers own their arithmetic and describe it in ``identity``. They must not
+    retain or mutate borrowed states. No callback is installed on the model between calls.
+
+    ``pooling`` (token runs only) pools the finished rows of each sequence over every attended token, as a
+    pooled ``HiddenTap`` does, so one value per sequence is kept instead of one per token; ``dtype`` converts
+    the finished rows first (float32 for float32 moments of a 16-bit reducer).
+    """
+
+    name: str
+    layers: tuple[int, ...]
+    begin: Callable[[], LayerAccumulator]
+    identity: Mapping[str, Any]
+    required_state_count: int | None = None
+    pooling: str | Sequence[str] | None = None
+    dtype: torch.dtype | None = None
+
+    def __post_init__(self) -> None:
+        if self.dtype is not None and self.dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
+            raise ValueError("A streaming tap dtype must be float16, bfloat16, float32, float64, or None.")
+        if self.pooling is not None:
+            names = Pooler(self.pooling).names  # validates names and rejects duplicates
+            if "parti" in names:
+                raise ValueError(f"Tap {self.name!r} cannot pool with 'parti', which needs the attention graph.")
+            object.__setattr__(self, "pooling", names)
+        layers = tuple(self.layers)
+        if not layers or any(type(layer) is not int or layer < 0 for layer in layers):
+            raise ValueError("Streaming layers must be nonempty nonnegative integer indices.")
+        if tuple(sorted(set(layers))) != layers:
+            raise ValueError("Streaming layers must be distinct and ascending.")
+        object.__setattr__(self, "layers", layers)
+        if self.required_state_count is not None and (
+            type(self.required_state_count) is not int or self.required_state_count <= layers[-1]
+        ):
+            raise ValueError("required_state_count must include every streamed layer.")
+        # Reuse the reducer identity validation and detached canonical copy.
+        checked = ReducedTap(self.name, layers[-1], self.begin, self.identity)
+        object.__setattr__(self, "identity", checked.identity)
+
+    @property
+    def layer(self) -> int:
+        """The deepest required state, for the existing early-stop plan."""
+        return self.layers[-1]
+
+
+Tap = HiddenTap | ReducedTap | StreamingTap | SparseResidueTap
+
+
+@dataclass(frozen=True, slots=True)
+class TapPlan:
+    """Validated taps, each layer resolved to a hidden-state index in ``0..n``."""
+
+    taps: tuple[Tap, ...]
+    layers: tuple[int, ...]
+
+    @property
+    def captured_layers(self) -> tuple[int, ...]:
+        """The distinct hidden states the forward pass must record, in ascending order."""
+
+        return tuple(sorted({
+            layer for tap, layer in zip(self.taps, self.layers, strict=True)
+            if not isinstance(tap, StreamingTap)
+        }))
+
+    @property
+    def streamed_layers(self) -> tuple[int, ...]:
+        return tuple(sorted({
+            layer for tap in self.taps if isinstance(tap, StreamingTap) for layer in tap.layers
+        }))
+
+    @property
+    def deepest_layer(self) -> int:
+        """The hidden state after which the forward pass stops."""
+
+        return max(self.layers)
+
+    @property
+    def pooling_names(self) -> frozenset[str]:
+        """Every pooler name the plan's hidden taps request."""
+
+        return frozenset(
+            name
+            for tap in self.taps
+            if isinstance(tap, HiddenTap) and tap.pooling is not None
+            for name in tap.pooling
+        )
+
+    def identity(self) -> list[dict[str, Any]]:
+        """Every tap in plan order, as the run fingerprint and metadata record it."""
+
+        described: list[dict[str, Any]] = []
+        for tap, layer in zip(self.taps, self.layers, strict=True):
+            if isinstance(tap, HiddenTap):
+                pooling = None if tap.pooling is None else list(tap.pooling)
+                described.append(
+                    {
+                        "name": tap.name,
+                        "kind": "hidden",
+                        "layer": layer,
+                        "pooling": pooling,
+                        "dtype": (
+                            str(tap.dtype).removeprefix("torch.") if tap.dtype is not None else None
+                        ),
+                    }
+                )
+            elif isinstance(tap, StreamingTap):
+                record = {
+                    "name": tap.name, "kind": "streaming", "layers": list(tap.layers),
+                    "identity": dict(tap.identity),
+                    "required_state_count": tap.required_state_count,
+                }
+                if tap.pooling is not None:  # absent for a per-token tap, so its existing fingerprint holds
+                    record.update(pooling=list(tap.pooling),
+                                  dtype=None if tap.dtype is None else str(tap.dtype).removeprefix("torch."))
+                described.append(record)
+            elif isinstance(tap, SparseResidueTap):
+                described.append({
+                    "name": tap.name, "kind": "sparse_residue", "layer": layer,
+                    "identity": dict(tap.identity), "codebook_size": tap.codebook_size,
+                    "sparse_count": tap.sparse_count,
+                })
+            else:
+                described.append(
+                    {
+                        "name": tap.name,
+                        "kind": "reduced",
+                        "layer": layer,
+                        "identity": dict(tap.identity),
+                    }
+                )
+        return described
+
+
+def plan_taps(taps: object, state_count: int) -> TapPlan:
+    """Validate ``taps`` against a model that exposes ``state_count`` hidden states.
+
+    ``taps`` is the caller's ``embed_dataset`` argument, checked here rather than trusted.
+    """
+
+    if isinstance(taps, (str, bytes)) or not isinstance(taps, Sequence):
+        raise TypeError(
+            "taps must be a sequence of HiddenTap, ReducedTap, StreamingTap "
+            "or SparseResidueTap values."
+        )
+    if not taps:
+        raise ValueError("taps must contain at least one tap.")
+    checked: list[Tap] = []
+    for tap in taps:
+        if not isinstance(tap, (HiddenTap, ReducedTap, StreamingTap, SparseResidueTap)):
+            raise TypeError(
+                "taps must contain HiddenTap, ReducedTap, StreamingTap or SparseResidueTap values; "
+                f"found {type(tap).__name__}."
+            )
+        checked.append(tap)
+    names = [tap.name for tap in checked]
+    repeated = sorted({name for name in names if names.count(name) > 1})
+    if repeated:
+        raise ValueError(f"Tap names must be unique; repeated: {repeated}.")
+    layers: list[int] = []
+    for tap in checked:
+        if isinstance(tap, StreamingTap) and tap.required_state_count not in (None, state_count):
+            raise ValueError(
+                f"Tap {tap.name!r} requires {tap.required_state_count} hidden states, "
+                f"not {state_count}."
+            )
+        if not -state_count <= tap.layer < state_count:
+            raise ValueError(
+                f"Tap {tap.name!r} names layer {tap.layer}, outside this model's hidden states "
+                f"{-state_count}..{state_count - 1}. Index i is the input to block i, and "
+                f"{state_count - 1} or -1 is the final normalized state."
+            )
+        layers.append(tap.layer % state_count)
+    return TapPlan(taps=tuple(checked), layers=tuple(layers))
+
+
+def _require_name_and_layer(name: object, layer: object) -> None:
+    if not isinstance(name, str) or not name:
+        raise ValueError("A tap name must be a non-empty string.")
+    if not isinstance(layer, int) or isinstance(layer, bool):
+        raise TypeError(f"Tap {name!r} layer must be an integer hidden-state index.")
+
+
+def _require_plain_data(value: object, where: str) -> None:
+    """Reject content whose serialized form could differ between runs of one reducer."""
+
+    if value is None or isinstance(value, (str, bool, int)):
+        return
+    if isinstance(value, float):
+        if not math.isfinite(value):
+            raise ValueError(f"The {where} holds a non-finite number.")
+        return
+    if isinstance(value, Mapping):
+        for key, item in value.items():
+            if not isinstance(key, str):
+                raise TypeError(f"The {where} has a non-string key {key!r}.")
+            _require_plain_data(item, where)
+        return
+    if isinstance(value, (list, tuple)):
+        for item in value:
+            _require_plain_data(item, where)
+        return
+    raise TypeError(
+        f"The {where} holds a {type(value).__name__}. An identity holds only strings, numbers, "
+        "booleans, None, lists, and string-keyed mappings, so its fingerprint is stable."
+    )
+
+
+__all__ = [
+    "HiddenTap",
+    "LayerAccumulator",
+    "ReducedTap",
+    "RowSelection",
+    "SparseResidueTap",
+    "StreamingTap",
+    "Tap",
+    "TapBatch",
+    "TapPlan",
+    "plan_taps",
+]

fastplms/embeddings/token_batches.py

@@ -0,0 +1,391 @@
+"""Embed canonical proteins with CLS and EOS kept: one pass per token-budget batch, no host stall, pinned outputs.
+
+A canonical run keeps the special tokens. A protein of l residues, after the N-terminal crop, is l + 2 rows in
+every per-token stream: row 0 CLS, rows 1..l residues, row l + 1 EOS. Padding is the only masked position, so a
+pooled vector averages all l + 2 rows. A legacy residue-only store has l rows (b, l, d) instead; the two never
+mix, because a v2 descriptor says ``special_tokens: kept``.
+
+The executor never reads a device value while it builds or runs a batch. Lengths are known on the host, so the
+attention mask, the row selection, and every offset come from host arrays, and the finite check is one device
+flag read after the outputs land. That keeps the device queue full while the previous batch is written.
+
+A run with a ``BatchGeometry`` gives every sequence one batch shape, whatever its companions: its l + 2 tokens round
+up to a bucket of T columns, and every batch of that bucket holds exactly ``rows(T)`` sequences padded to T. A
+GEMM's reduction order follows its shape, so a fixed shape per sequence is what makes a stored row independent of
+the batch that made it. Without a geometry, batches follow the token budget and pad to their longest member.
+
+Symbols: b sequences of a batch; l residues of one sequence after the crop; m = max(l) + 2 padded token columns,
+or the bucket T of a geometry batch; n = sum(l_i + 2) attended token rows; d hidden width; c SAE codebook; k SAE
+codes kept per token; w stored columns of a stream.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+from torch import Tensor
+
+from .pooling import pool_token_rows
+from .taps import HiddenTap, ReducedTap, RowSelection, SparseResidueTap, StreamingTap, TapBatch, TapPlan
+from .tokens import ResidueVocabulary
+from ..features.async_writer import PackedBatch
+from ..features.layouts import INDEX_DTYPE
+
+
+# CLS before the residues and EOS after them.
+SPECIAL_TOKEN_ROWS = 2
+# The N-terminal crop in residues: with CLS and EOS it fills a 2,048-token context. The one place the engine names it.
+CANONICAL_MAX_RESIDUES = 2046
+# A geometry run's batch algorithm and its rule for a bucket's last, partial batch.
+GEOMETRY_ALGORITHM = "bucketed_fixed_shape_v1"
+PARTIAL_BATCH_POLICY = "repeat_last_sequence_discard_duplicate_outputs_v1"
+
+
+@dataclass(frozen=True, slots=True)
+class BatchGeometry:
+    """Fixed batch shapes, so a sequence meets the same kernels in whatever batch it runs.
+
+    A sequence of l residues after the crop needs l + 2 token columns and runs in the bucket
+    T = ``bucket_tokens`` * ceil((l + 2) / ``bucket_tokens``), at most ``max_columns``. Every batch of bucket T holds
+    exactly ``rows(T)`` sequences, as many as ``token_budget`` padded token rows hold and at most ``max_rows``, each
+    padded to T columns; a bucket's last batch repeats its last sequence to fill and discards the copies' outputs.
+    """
+
+    bucket_tokens: int
+    token_budget: int  # padded token rows one batch may hold
+    max_rows: int  # sequences one batch may hold
+    max_columns: int = CANONICAL_MAX_RESIDUES + SPECIAL_TOKEN_ROWS
+
+    def __post_init__(self) -> None:
+        for name in ("bucket_tokens", "token_budget", "max_rows", "max_columns"):
+            value = getattr(self, name)
+            if type(value) is not int or value < 1:
+                raise ValueError(f"BatchGeometry.{name} must be a positive integer; received {value!r}.")
+        if self.max_columns % self.bucket_tokens:
+            raise ValueError("BatchGeometry.max_columns must be a multiple of bucket_tokens.")
+        if self.token_budget < self.max_columns:
+            raise ValueError("BatchGeometry.token_budget must hold one sequence of max_columns tokens.")
+
+    def columns(self, residues: int) -> int:
+        """The bucket T of a sequence of ``residues`` after the crop: its l + 2 tokens rounded up to the bucket width."""
+        tokens = residues + SPECIAL_TOKEN_ROWS
+        if residues < 1 or tokens > self.max_columns:
+            raise ValueError(
+                f"A sequence of {residues} residues needs {tokens} token columns; this geometry holds 3 to {self.max_columns}."
+            )
+        return -(-tokens // self.bucket_tokens) * self.bucket_tokens
+
+    def rows(self, columns: int) -> int:
+        """Sequences in every batch of bucket ``columns``: as many as the token budget holds, at most ``max_rows``."""
+        if columns % self.bucket_tokens or not 0 < columns <= self.max_columns:
+            raise ValueError(f"{columns} is not a bucket of this geometry.")
+        return min(self.max_rows, self.token_budget // columns)
+
+    def shapes(self) -> dict[int, int]:
+        """Every bucket T with the row count rows(T) of its batches."""
+        return {
+            columns: self.rows(columns)
+            for columns in range(self.bucket_tokens, self.max_columns + 1, self.bucket_tokens)
+        }
+
+    def describe(self) -> dict[str, int | str]:
+        """The geometry as a run record and a feature contract state it."""
+        return {
+            "algorithm": GEOMETRY_ALGORITHM, "bucket_tokens": self.bucket_tokens, "token_budget": self.token_budget,
+            "max_rows": self.max_rows, "max_columns": self.max_columns, "partial_batch": PARTIAL_BATCH_POLICY,
+        }
+
+
+def plan_geometry_batches(
+    lengths: Sequence[int], digests: Sequence[str], geometry: BatchGeometry,
+) -> Iterator[tuple[int, ...]]:
+    """Batches of one bucket each: the longest bucket first, members in row-key order, ``rows(T)`` to a batch.
+
+    ``lengths`` holds the residues l of each sequence after the crop and ``digests`` their row keys. A bucket's last
+    batch may hold fewer sequences; the executor fills it to ``rows(T)``. The longest bucket first makes an
+    out-of-memory failure show on the first batch, and key order makes the plan independent of the input's order.
+    Yields indices into ``lengths``.
+    """
+    if len(lengths) != len(digests):
+        raise ValueError("plan_geometry_batches needs one row key per length.")
+    buckets: dict[int, list[int]] = {}
+    for index, residues in enumerate(lengths):
+        buckets.setdefault(geometry.columns(residues), []).append(index)
+    for columns in sorted(buckets, reverse=True):
+        members = sorted(buckets[columns], key=lambda index: digests[index])
+        size = geometry.rows(columns)
+        for start in range(0, len(members), size):
+            yield tuple(members[start:start + size])
+
+
+def plan_token_batches(
+    lengths: Sequence[int], *, max_sequences: int, max_tokens: int, window: int,
+) -> Iterator[tuple[int, ...]]:
+    """Group sequences into batches of similar length under a padded-token budget.
+
+    ``lengths`` holds the residues l of each sequence after the crop. Sequences are sorted longest first
+    inside each window of ``window`` consecutive sequences, so a batch pads to its first member: it holds
+    at most ``max_sequences`` sequences and ``b * (l_first + 2) <= max_tokens`` padded token rows. Longest
+    first also makes an out-of-memory failure show on the first batch. Yields indices into ``lengths``.
+    """
+    if min(max_sequences, max_tokens, window) < 1:
+        raise ValueError("max_sequences, max_tokens and window must be positive.")
+    for start in range(0, len(lengths), window):
+        order = sorted(range(start, min(start + window, len(lengths))), key=lambda index: (-lengths[index], index))
+        batch: list[int] = []
+        for index in order:
+            if lengths[index] + SPECIAL_TOKEN_ROWS > max_tokens:
+                raise ValueError(
+                    f"A sequence of {lengths[index]} residues needs {lengths[index] + SPECIAL_TOKEN_ROWS} token rows, "
+                    f"more than max_tokens={max_tokens}."
+                )
+            padded_rows = (len(batch) + 1) * (lengths[batch[0]] + SPECIAL_TOKEN_ROWS) if batch else 0
+            if batch and (len(batch) + 1 > max_sequences or padded_rows > max_tokens):
+                yield tuple(batch)
+                batch = []
+            batch.append(index)
+        if batch:
+            yield tuple(batch)
+
+
+@dataclass(frozen=True, slots=True)
+class HostBatch:
+    """The integer arrays of one batch, built on the host from known lengths."""
+
+    input_ids: np.ndarray  # (b, m) int64, CLS, residue ids, EOS, then padding
+    rows: np.ndarray  # (b,) int64, l_i + 2 attended tokens of each sequence
+    flat_index: np.ndarray  # (n,) int64, each attended token's position in the row-major (b * m) grid
+    owner: np.ndarray  # (n,) int64, the sequence of each attended token
+
+
+def build_host_batch(vocabulary: ResidueVocabulary, texts: Sequence[str], *, columns: int | None = None) -> HostBatch:
+    """Token ids, lengths, and the row selection of ``texts`` (already cropped), with no device value.
+
+    ``columns`` pads every sequence to that many token columns (a geometry batch's bucket T); None pads to the
+    longest sequence.
+    """
+    encoded = [vocabulary.encode(text) for text in texts]  # b arrays of (l_i + 2,)
+    rows = np.fromiter((len(ids) for ids in encoded), dtype=np.int64, count=len(encoded))  # (b,)
+    longest = int(rows.max())
+    if columns is not None and columns < longest:
+        raise ValueError(f"A batch padded to {columns} columns cannot hold a sequence of {longest} tokens.")
+    m = longest if columns is None else columns
+    input_ids = np.full((len(encoded), m), vocabulary.pad_id, dtype=np.int64)  # (b, m)
+    for index, ids in enumerate(encoded):
+        input_ids[index, : len(ids)] = ids
+    owner = np.repeat(np.arange(len(encoded), dtype=np.int64), rows)  # (n,)
+    starts = np.cumsum(rows) - rows  # (b,) first packed row of each sequence
+    within = np.arange(int(rows.sum()), dtype=np.int64) - np.repeat(starts, rows)  # (n,) token index in its sequence
+    return HostBatch(input_ids, rows, owner * m + within, owner)
+
+
+def _to_device(arrays: Sequence[np.ndarray], device: torch.device) -> list[Tensor]:
+    """Copy several host arrays to the device in one transfer from pinned memory, without blocking the host."""
+    # arrays: (b, m), (b,), (n,), (n,) int64 for the executor's batch; each comes back with its own shape.
+    flat = np.concatenate([array.reshape(-1) for array in arrays])  # (total,) int64
+    if device.type == "cuda":
+        staging = torch.empty(flat.shape[0], dtype=torch.int64, pin_memory=True)  # (total,)
+        staging.numpy()[:] = flat
+        moved = staging.to(device, non_blocking=True)  # (total,)
+    else:
+        moved = torch.from_numpy(flat).to(device)  # (total,)
+    pieces, cursor = [], 0
+    for array in arrays:
+        pieces.append(moved[cursor : cursor + array.size].view(array.shape))
+        cursor += array.size
+    return pieces  # views shaped like arrays, e.g. (b, m), (b,), (n,), (n,), of one device buffer
+
+
+class TokenTapExecutor:
+    """Run a tap plan over canonical proteins and hand each batch to the writer as pinned host tensors.
+
+    ``plan`` holds the taps. Every tap sees all attended tokens, CLS and EOS included, so a hidden tap keeps
+    ``(n, d)`` rows and a pooled tap averages l + 2 rows. ``max_residues`` is the N-terminal crop. ``dtype`` is
+    the run dtype a tap converts to unless it names its own, and None keeps the model's dtype. ``geometry`` runs
+    every batch at its bucket's fixed shape (``BatchGeometry``); ``fixed_batch_size`` fills every batch to that
+    many sequences but pads it to its longest member. A run uses one of the two at most.
+    """
+
+    def __init__(
+        self, model: Any, plan: TapPlan, *, vocabulary: ResidueVocabulary, max_residues: int | None,
+        dtype: torch.dtype | None,
+        fixed_batch_size: int | None = None,
+        geometry: BatchGeometry | None = None,
+    ) -> None:
+        if getattr(model, "embedding_tap_support", False) is not True:
+            raise ValueError(f"{type(model).__name__} does not support one-pass taps.")
+        if any(isinstance(tap, StreamingTap) for tap in plan.taps) and getattr(
+            model, "embedding_streaming_tap_support", False
+        ) is not True:
+            raise ValueError("This model does not support streaming hidden-state taps.")
+        if max_residues is not None and max_residues < 1:
+            raise ValueError("max_residues must be positive.")
+        if fixed_batch_size is not None and (type(fixed_batch_size) is not int or fixed_batch_size < 1):
+            raise ValueError("fixed_batch_size must be a positive integer.")
+        if geometry is not None:
+            if fixed_batch_size is not None:
+                raise ValueError("A run takes its batch shapes from a geometry or a fixed batch size, not both.")
+            if max_residues is None or max_residues + SPECIAL_TOKEN_ROWS > geometry.max_columns:
+                raise ValueError("A geometry run needs a crop whose l + 2 tokens fit the geometry's widest bucket.")
+        self.model = model
+        self.plan = plan
+        self.vocabulary = vocabulary
+        self.max_residues = max_residues
+        self.dtype = dtype
+        self.fixed_batch_size = fixed_batch_size
+        self.geometry = geometry
+        self.device = next(model.parameters()).device
+
+    def crop(self, sequence: str) -> str:
+        """The N-terminal crop: the first ``max_residues`` residues, so l <= max_residues and l + 2 tokens."""
+        return sequence if self.max_residues is None else sequence[: self.max_residues]
+
+    def batch_shape(self, texts: Sequence[str]) -> tuple[int, int] | None:
+        """The (rows, columns) a geometry runs ``texts`` (already cropped) at; None without a geometry."""
+        if self.geometry is None:
+            return None
+        buckets = {self.geometry.columns(len(text)) for text in texts}
+        if len(buckets) != 1:
+            raise ValueError(f"A geometry batch holds sequences of one bucket; these fall in {sorted(buckets)}.")
+        (columns,) = buckets
+        rows = self.geometry.rows(columns)
+        if len(texts) > rows:
+            raise ValueError(f"Bucket {columns} runs {rows} sequences to a batch; received {len(texts)}.")
+        return rows, columns
+
+    def run_batch(self, sequences: Sequence[str], digests: Sequence[str]) -> PackedBatch:
+        """Embed one batch and start the copies of its outputs to the host; returns without waiting for them."""
+        count = len(sequences)
+        if not count or len(digests) != count:
+            raise ValueError("A batch needs at least one sequence and one digest per sequence.")
+        if self.fixed_batch_size is not None and count > self.fixed_batch_size:
+            raise ValueError("The input exceeds the fixed physical batch size.")
+        texts = [self.crop(sequence) for sequence in sequences]
+        shape = self.batch_shape(texts)
+        columns = None
+        if shape is not None:
+            rows, columns = shape
+            texts += [texts[-1]] * (rows - count)
+        elif self.fixed_batch_size is not None:
+            texts += [texts[-1]] * (self.fixed_batch_size - count)
+        host = build_host_batch(self.vocabulary, texts, columns=columns)
+        input_ids, rows, flat_index, owner = _to_device(
+            (host.input_ids, host.rows, host.flat_index, host.owner), self.device,
+        )  # (b, m), (b,), (n,), (n,)
+        m = input_ids.shape[1]
+        # (b, m): true on CLS, residues and EOS; padding is the only masked position.
+        token_mask = torch.arange(m, device=self.device).unsqueeze(0) < rows.unsqueeze(1)  # (b, m)
+        selection = RowSelection(flat_index, owner, tuple(int(count) for count in host.rows), rows)
+        cache: dict[Any, Any] = {}
+
+        def batch_for(X: Tensor) -> TapBatch:
+            # X: (b, m, d) one layer's hidden state; padding columns stay in X and are masked by token_mask (b, m).
+            return TapBatch(X=X, token_mask=token_mask, residue_mask=token_mask, rows=selection, cache=cache)
+
+        streaming = tuple(tap for tap in self.plan.taps if isinstance(tap, StreamingTap))
+        accumulators = {tap.name: tap.begin() for tap in streaming}
+        stream_layers = self.plan.streamed_layers
+
+        def consume(layer: int, X: Tensor) -> None:
+            # X: (b, m, d), borrowed until this callback returns
+            borrowed = batch_for(X)
+            for tap in streaming:
+                if layer in tap.layers:
+                    accumulators[tap.name].update(layer, borrowed)
+
+        states = self.model._embed_taps(
+            input_ids, token_mask, self.plan.captured_layers, stream_layers=stream_layers,
+            state_consumer=consume if streaming else None, assume_valid_mask=True,
+        )  # {layer: (b, m, d)}
+        outputs: dict[str, dict[str, Tensor]] = {}
+        for tap, layer in zip(self.plan.taps, self.plan.layers, strict=True):
+            if isinstance(tap, StreamingTap):
+                Y = accumulators[tap.name].finish()  # (b, m, w)
+                if tap.dtype is not None:
+                    Y = Y.to(tap.dtype)  # (b, m, w)
+                if tap.pooling is None:
+                    outputs[tap.name] = {"values": selection.gather(Y)}  # (n, w)
+                else:
+                    outputs[tap.name] = {"values": pool_token_rows(Y, token_mask, tap.pooling)}  # (b, p * w)
+                continue
+            dtype = tap.dtype if isinstance(tap, HiddenTap) and tap.dtype is not None else self.dtype
+            X = states[layer]  # (b, m, d)
+            if dtype is not None:
+                X = X.to(dtype)  # (b, m, d)
+            if isinstance(tap, SparseResidueTap):
+                if tap.reduce_packed is None:
+                    raise ValueError(f"Tap {tap.name!r} has no packed reducer; it cannot keep special tokens.")
+                packed = tap.reduce_packed(batch_for(X))  # indices, values (n, k)
+                outputs[tap.name] = {  # indices and values: (n, k)
+                    "indices": packed.indices.to(INDEX_DTYPE),
+                    "values": packed.values,
+                }
+            elif isinstance(tap, ReducedTap):
+                outputs[tap.name] = {"values": tap.reduce(batch_for(X))}  # (b, w)
+            elif tap.pooling is None:
+                outputs[tap.name] = {"values": selection.gather(X)}  # (n, d)
+            else:
+                outputs[tap.name] = {"values": pool_token_rows(X, token_mask, tap.pooling)}  # (b, p * d)
+        if len(texts) != count:
+            token_rows = int(host.rows[:count].sum())
+            for tap in self.plan.taps:
+                per_token = isinstance(tap, SparseResidueTap) or (
+                    isinstance(tap, (HiddenTap, StreamingTap)) and tap.pooling is None)
+                retained = token_rows if per_token else count
+                outputs[tap.name] = {name: tensor[:retained] for name, tensor in outputs[tap.name].items()}
+        return self._deliver(outputs, sequences, digests, tuple(int(rows) for rows in host.rows[:count]))
+
+    def _deliver(
+        self, outputs: Mapping[str, Mapping[str, Tensor]], sequences: Sequence[str], digests: Sequence[str],
+        rows: tuple[int, ...],
+    ) -> PackedBatch:
+        """Start the device-to-host copies, record one event after them, and check finiteness on the device."""
+        # outputs: (b, w) pooled, (n, w) per token row, (n, k) top-k indices, per tap; n = sum(l_i + 2).
+        finite = [torch.isfinite(tensor).all() for group in outputs.values() for tensor in group.values()
+                  if tensor.is_floating_point()]  # () per floating tensor
+        everything_finite = torch.stack(finite).all() if finite else torch.ones((), dtype=torch.bool, device=self.device)  # ()
+        if self.device.type != "cuda":
+            hosts = {name: dict(group) for name, group in outputs.items()}
+            verdict = everything_finite
+
+            def wait() -> None:
+                if not bool(verdict):
+                    raise ValueError("A tap produced a non-finite value.")
+        else:
+            hosts = {name: {key: _pinned_copy(tensor) for key, tensor in group.items()} for name, group in outputs.items()}
+            flag = _pinned_copy(everything_finite)
+            event = torch.cuda.Event()
+            event.record()
+
+            def wait() -> None:
+                event.synchronize()
+                if not bool(flag):
+                    raise ValueError("A tap produced a non-finite value.")
+
+        return PackedBatch(tuple(sequences), tuple(digests), rows, hosts, wait)
+
+
+def _pinned_copy(tensor: Tensor) -> Tensor:
+    """A pinned host tensor the copy of ``tensor`` is queued into on the current stream, without waiting."""
+    # tensor: (...) any shape; the host copy has the same shape and dtype.
+    host = torch.empty(tensor.shape, dtype=tensor.dtype, pin_memory=True)  # (...)
+    host.copy_(tensor.detach(), non_blocking=True)
+    return host  # (...) the shape of tensor
+
+
+__all__ = [
+    "CANONICAL_MAX_RESIDUES",
+    "GEOMETRY_ALGORITHM",
+    "PARTIAL_BATCH_POLICY",
+    "SPECIAL_TOKEN_ROWS",
+    "BatchGeometry",
+    "HostBatch",
+    "TokenTapExecutor",
+    "build_host_batch",
+    "plan_geometry_batches",
+    "plan_token_batches",
+]

fastplms/embeddings/token_runs.py

@@ -0,0 +1,252 @@
+"""Fill feature stores from canonical proteins with CLS and EOS kept, through the asynchronous writer.
+
+``embed_token_features`` is the fast path of ``embed_into_features(keep_special_tokens=True)``. It takes a
+protein inventory, embeds only the rows a stream lacks, and writes each tap's rows into the store of its key.
+The device loop (``TokenTapExecutor``) and the disk loop (``AsyncFeatureWriter``) run concurrently, joined by a
+bounded queue of pinned host buffers, so the model is not paused for hashing, compression, or fsync.
+
+Per batch, every per-token stream holds l + 2 rows per sequence (row 0 CLS, rows 1..l residues, row l + 1
+EOS) and every pooled stream averages or maximizes over those same l + 2 rows. A legacy residue-only store
+holds l rows; this path never writes one.
+
+Symbols: b sequences of a batch; l residues of a sequence after the N-terminal crop; n = sum(l_i + 2) token rows.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import torch
+
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Protocol
+
+from .batches import _temporary_eval
+from .pooling import POOLING_SEMANTICS_TOKENS
+from .taps import Tap, plan_taps
+from .token_batches import (
+    CANONICAL_MAX_RESIDUES, SPECIAL_TOKEN_ROWS, BatchGeometry, TokenTapExecutor, plan_geometry_batches, plan_token_batches,
+)
+from .tokens import ResidueVocabulary, check_canonical_text
+from ..features.async_writer import AsyncFeatureWriter
+from ..features.store import FeatureStore, SegmentReceipt, StoredFeature, sequence_digest
+
+
+GIB = 1024**3
+TOKEN_RUN_SCHEMA = "token_features_v1"
+BATCH_ALGORITHM = "token_budget_length_sorted_v1"
+# The token budget's defaults, for a run without a geometry.
+DEFAULT_MAX_SEQUENCES = 256
+DEFAULT_MAX_TOKENS = 32768
+DEFAULT_WINDOW = 65536
+
+
+class TokenFeatureContract(Protocol):
+    """What a token run asks of a scientific contract: validate it, name each row, and recheck before a commit."""
+
+    def validate(
+        self, model: Any, sequences: Sequence[str], features: Mapping[str, StoredFeature],
+        taps: Sequence[Tap], options: Mapping[str, Any],
+    ) -> None: ...
+
+    def validate_cached(self, name: str, store: FeatureStore, sequences: Sequence[str]) -> None: ...
+
+    def row_identities(
+        self, name: str, sequences: Sequence[str], digests: Sequence[str],
+    ) -> Sequence[Mapping[str, Any]]: ...
+
+    def check_row_keys(self, sequences: Sequence[str], digests: Sequence[str]) -> None:
+        """Raise unless each digest is the key the contract holds for its sequence."""
+
+    def before_commit(self) -> None: ...
+
+
+def distinct_with_digests(
+    sequences: Sequence[str], digests: Sequence[str] | None,
+) -> tuple[list[str], list[str]]:
+    """Distinct sequences in first-seen order and their row keys, hashing each exact text once.
+
+    ``digests`` may supply the keys when the caller already verified them against the text.
+    """
+    if digests is not None and len(digests) != len(sequences):
+        raise ValueError("digests must hold one key per sequence.")
+    seen: set[str] = set()
+    ordered: list[str] = []
+    keys: list[str] = []
+    for position, sequence in enumerate(sequences):
+        if sequence in seen:
+            continue
+        seen.add(sequence)
+        ordered.append(sequence)
+        keys.append(sequence_digest(sequence) if digests is None else digests[position])
+    return ordered, keys
+
+
+def embed_token_features(
+    model: Any,
+    sequences: Sequence[str],
+    root: str | Path,
+    features: Mapping[str, StoredFeature],
+    *,
+    taps: Sequence[Tap],
+    contract: TokenFeatureContract | None = None,
+    digests: Sequence[str] | None = None,
+    metadata: Mapping[str, Any] | None = None,
+    max_residues: int | None = CANONICAL_MAX_RESIDUES,
+    max_sequences: int | None = None,
+    max_tokens: int | None = None,
+    window: int | None = None,
+    dtype: torch.dtype | None = None,
+    fixed_batch_size: int | None = None,
+    geometry: BatchGeometry | None = None,
+    part_bytes: int = GIB,
+    segment_bytes: int = 8 * GIB,
+    queue_bytes: int = 4 * GIB,
+    workers: int = 4,
+    verify_cached: bool = True,
+    model_state_fingerprint: str | None = None,
+    progress: Callable[[int], None] | None = None,
+    on_plan: Callable[[int], None] | None = None,
+    batch_watch: Callable[[Iterator[tuple[int, ...]]], Iterator[tuple[int, ...]]] | None = None,
+) -> dict[str, tuple[SegmentReceipt, ...]]:
+    """Fill each named feature under ``root`` from one pass over the sequences it lacks, special tokens kept.
+
+    ``features`` maps a tap name to its feature and names every tap. ``max_residues`` is the N-terminal
+    crop (``CANONICAL_MAX_RESIDUES`` keeps a sequence within 2048 tokens). Batches hold at most ``max_sequences`` sequences and
+    ``max_tokens`` padded token rows, sorted by length inside windows of ``window`` sequences (256, 32768 and 65536 when
+    None). A ``geometry`` instead runs every sequence at its bucket's fixed shape (``BatchGeometry``), whatever batch it
+    falls in, and takes none of those three nor ``fixed_batch_size``. Parts are
+    about ``part_bytes``, a segment commits every ``segment_bytes`` across all streams, and at most
+    ``queue_bytes`` of finished rows wait for the writer. ``digests`` are the rows' SHA-256 keys when the
+    caller already has them; otherwise each text is hashed once here. ``verify_cached=False`` skips the
+    contract's re-read of rows already stored, which a resumed run of a large store does separately.
+    ``progress`` receives the sequences of each batch once its rows are packed for writing, and ``on_plan`` the number of
+    sequences this run will embed (fewer than ``sequences`` on a resume), before the first batch.
+
+    Returns the committed segments of each stream that gained rows, oldest first; an empty result means the
+    model never ran. A killed run commits whole segments only, so a rerun resumes from the last one.
+    """
+    names = {tap.name for tap in taps}
+    if geometry is not None:
+        if any(value is not None for value in (max_sequences, max_tokens, window, fixed_batch_size)):
+            raise ValueError(
+                "A geometry run takes its batch shapes from the geometry; pass no max_sequences, max_tokens, window "
+                "or fixed_batch_size."
+            )
+        if max_residues is None or max_residues + SPECIAL_TOKEN_ROWS > geometry.max_columns:
+            raise ValueError("A geometry run needs a crop whose l + 2 tokens fit the geometry's widest bucket.")
+    else:
+        max_sequences = DEFAULT_MAX_SEQUENCES if max_sequences is None else max_sequences
+        max_tokens = DEFAULT_MAX_TOKENS if max_tokens is None else max_tokens
+        window = DEFAULT_WINDOW if window is None else window
+    if fixed_batch_size is not None:
+        if fixed_batch_size != max_sequences or max_residues is None or max_tokens < fixed_batch_size * (max_residues + 2):
+            raise ValueError("Fixed batches require matching max_sequences and a token budget covering the full cropped context.")
+    if set(features) != names:
+        raise ValueError(
+            "features must name exactly the taps this run takes.\n"
+            f"  taps:     {sorted(names)}\n  features: {sorted(features)}"
+        )
+    if any(spec.positions for spec in features.values()):
+        raise ValueError("A token run stores no argmax positions.")
+    if contract is None and any(spec.descriptor.get("schema") is not None for spec in features.values()):
+        raise ValueError("A descriptor that carries a schema (feature_spec_v1, v2 or v3) requires its contract.")
+    if getattr(contract, "keep_special_tokens", None) is False:
+        raise ValueError("A token run needs a contract captured with the special tokens kept.")
+    ordered, keys = distinct_with_digests(sequences, digests)
+    if not ordered:
+        raise ValueError("embed_token_features needs at least one sequence.")
+    # A row is keyed by the hash of its text, so the text must be the normalized one (uppercase, no whitespace):
+    # a second spelling of a protein would otherwise become a second row. Checked for all before any forward.
+    for sequence in ordered:
+        check_canonical_text(sequence)
+    if contract is not None:
+        contract.check_row_keys(ordered, keys)  # a dict lookup per row: a caller's digest is never trusted
+        # The contract compares these with the options it measured; its per-sequence checks are the caller's.
+        extraction_options: dict[str, Any] = {"max_length": max_residues, "truncate": True, "dtype": dtype}
+        if geometry is not None:
+            extraction_options["geometry"] = geometry.describe()
+        else:
+            extraction_options.update(batch_size=max_sequences, batch_window_size=window, max_tokens_per_batch=max_tokens)
+        if fixed_batch_size is not None:
+            extraction_options["fixed_batch_size"] = fixed_batch_size
+        contract.validate(model, (), features, taps, extraction_options)
+    stores = {name: FeatureStore.open(root, spec, deep_verify=False) for name, spec in features.items()}
+    # Each stream's missing rows come from one index pass over the keys, never a second hash of the text.
+    wanted = {name: frozenset(set(keys) - store.present_digests(keys)) for name, store in stores.items()}
+    if verify_cached and contract is not None:
+        for name, store in stores.items():
+            cached = [sequence for sequence, key in zip(ordered, keys, strict=True) if key not in wanted[name]]
+            if cached:
+                contract.validate_cached(name, store, cached)
+    chosen = [position for position, key in enumerate(keys) if any(key in group for group in wanted.values())]
+    if not chosen:
+        return {}
+    texts = [ordered[position] for position in chosen]
+    text_keys = [keys[position] for position in chosen]
+    if on_plan is not None:
+        on_plan(len(texts))
+    plan = plan_taps(list(taps), int(model.embedding_tap_state_count))
+    executor = TokenTapExecutor(
+        model, plan, vocabulary=ResidueVocabulary(model.tokenizer), max_residues=max_residues, dtype=dtype,
+        fixed_batch_size=fixed_batch_size, geometry=geometry,
+    )
+    lengths = [len(executor.crop(text)) for text in texts]  # (count,) residues l after the crop
+    batch_policy: dict[str, Any] = dict(geometry.describe()) if geometry is not None else {
+        "algorithm": BATCH_ALGORITHM, "max_sequences": max_sequences, "max_tokens": max_tokens, "window": window,
+    }
+    run = {
+        "schema": TOKEN_RUN_SCHEMA, "special_tokens": "kept", "max_residues": max_residues,
+        "batch_policy": batch_policy,
+        "pooling_semantics": dict(POOLING_SEMANTICS_TOKENS), "model_state_fingerprint": model_state_fingerprint,
+        "storage_policy": {"max_part_bytes": part_bytes, "segment_bytes": segment_bytes},
+    }
+    if fixed_batch_size is not None:
+        run["batch_policy"].update(algorithm="fixed_rows_duplicate_pad_v1", fixed_batch_size=fixed_batch_size)
+    # Each stream's own missing rows name the segment. A kill after one stream committed leaves that stream
+    # with fewer missing rows, so the rerun gets new segment names and never collides with the committed one.
+    wanted_digests = {
+        features[name].key: hashlib.sha256("".join(sorted(group)).encode("ascii")).hexdigest()
+        for name, group in wanted.items()
+    }
+    fingerprint = hashlib.sha256(json.dumps({"run": run, "wanted": wanted_digests}, sort_keys=True)
+                                 .encode("utf-8")).hexdigest()[:32]
+
+    def identities(stream: str, batch_texts: Sequence[str], batch_keys: Sequence[str]) -> Sequence[Mapping[str, Any]]:
+        if contract is None:
+            return [{} for _ in batch_texts]
+        # Identities bind the original text, whose length the crop policy turns into l; never the cropped text.
+        return contract.row_identities(stream, batch_texts, batch_keys)
+
+    writer = AsyncFeatureWriter(
+        stores, fingerprint=fingerprint, metadata={**dict(metadata or {}), **run}, wanted=wanted,
+        row_records=identities, part_bytes=part_bytes, segment_bytes=segment_bytes, queue_bytes=queue_bytes,
+        before_commit=None if contract is None else contract.before_commit, workers=workers, progress=progress,
+    )
+    try:
+        with _inference(model):
+            if geometry is not None:
+                batches = plan_geometry_batches(lengths, text_keys, geometry)
+            else:
+                batches = plan_token_batches(lengths, max_sequences=max_sequences, max_tokens=max_tokens, window=window)
+            if batch_watch is not None:
+                batches = batch_watch(batches)
+            for members in batches:
+                writer.submit(executor.run_batch([texts[i] for i in members], [text_keys[i] for i in members]))
+    except BaseException:
+        writer.abort()
+        raise
+    receipts = writer.close()
+    return {name: tuple(group) for name, group in receipts.items() if group}
+
+
+@contextmanager
+def _inference(model: Any) -> Iterator[None]:
+    """Evaluation mode and inference mode for the whole loop, restored afterwards."""
+    with _temporary_eval(model), torch.inference_mode():
+        yield
+
+
+__all__ = ["BATCH_ALGORITHM", "TOKEN_RUN_SCHEMA", "TokenFeatureContract", "distinct_with_digests", "embed_token_features"]

fastplms/embeddings/tokens.py

@@ -0,0 +1,74 @@
+"""Tokenize canonical proteins once, as lookup-table rows, and check the table against the tokenizer.
+
+A canonical sequence is already normalized (uppercase ASCII letters), and an ESM tokenizer maps each
+residue letter to one id, so a protein of `l` residues is `l + 2` ids: CLS, the residues, EOS. The
+vocabulary builds a 256-entry table from the tokenizer once, then encodes a sequence by one array
+lookup instead of a Python call per residue. `verify` holds the table to the tokenizer itself, so
+the ids the model sees are the ids the tokenizer would have produced.
+
+Symbols: `l` residues of one sequence, so `l + 2` token ids (row 0 CLS, rows 1..l residues, row `l + 1` EOS).
+"""
+
+from __future__ import annotations
+
+import string
+import numpy as np
+
+from typing import Any
+from numpy.typing import NDArray
+
+
+UNKNOWN_ID = -1
+_LETTERS = string.ascii_uppercase
+
+
+def check_canonical_text(sequence: str) -> None:
+    """Reject text that is not an already normalized uppercase protein; this never normalizes."""
+    if not sequence or not sequence.isascii() or not sequence.isalpha() or not sequence.isupper():
+        raise ValueError("Canonical feature input must be an already normalized uppercase protein.")
+
+
+class ResidueVocabulary:
+    """Residue letter to token id, built once from a tokenizer and verified against it."""
+
+    def __init__(self, tokenizer: Any) -> None:
+        special = set(tokenizer.all_special_ids)
+        vocabulary = tokenizer.get_vocab()
+        table = np.full(256, UNKNOWN_ID, dtype=np.int64)  # (256,) ASCII code to token id
+        for letter in _LETTERS:
+            token_id = vocabulary.get(letter)
+            if token_id is not None and token_id not in special:
+                table[ord(letter)] = token_id
+        self.table = table
+        self.cls_id = int(tokenizer.cls_token_id)
+        self.eos_id = int(tokenizer.eos_token_id)
+        self.pad_id = int(tokenizer.pad_token_id)
+        self.verify(tokenizer)
+
+    def verify(self, tokenizer: Any) -> None:
+        """Hold every mapped letter to the tokenizer: `CLS, letter, EOS` must be its encoding."""
+        for letter in _LETTERS:
+            token_id = int(self.table[ord(letter)])
+            if token_id == UNKNOWN_ID:
+                continue
+            encoded = tokenizer([letter], add_special_tokens=True)["input_ids"][0]
+            if list(encoded) != [self.cls_id, token_id, self.eos_id]:
+                raise ValueError(
+                    f"Residue {letter!r} encodes as {list(encoded)}, not the table's "
+                    f"{[self.cls_id, token_id, self.eos_id]}; the tokenizer is not one token per residue."
+                )
+
+    def encode(self, sequence: str) -> NDArray[np.int64]:
+        """Token ids `(l + 2,)` of one canonical sequence: CLS, one id per residue, EOS."""
+        check_canonical_text(sequence)
+        letters = np.frombuffer(sequence.encode("ascii"), dtype=np.uint8)  # (l,)
+        residues = self.table[letters]  # (l,) token ids, UNKNOWN_ID where the tokenizer lacks the letter
+        if bool((residues == UNKNOWN_ID).any()):
+            raise ValueError("A canonical residue has no non-special tokenizer representation.")
+        ids = np.empty(len(letters) + 2, dtype=np.int64)  # (l+2,)
+        ids[0], ids[-1] = self.cls_id, self.eos_id
+        ids[1:-1] = residues
+        return ids  # (l+2,)
+
+
+__all__ = ["UNKNOWN_ID", "ResidueVocabulary", "check_canonical_text"]

fastplms/embeddings/types.py

@@ -7,6 +7,9 @@ from dataclasses import dataclass, field
 from typing import Any, Literal, overload
 from torch import Tensor
 
+from types import MappingProxyType
+from ..features.layouts import TopKRow
+
 
 @dataclass(frozen=True, slots=True)
 class EmbeddingInput:
@@ -38,7 +41,7 @@ class LazyTensorReference:
 
         if not isinstance(verify, bool):
             raise TypeError("verify must be a boolean.")
-        X = self._loader()  # self.shape
+        X = self._loader()  # (...), equal to self.shape
         if not isinstance(X, Tensor):
             raise TypeError(f"Stored tensor loader for {self.key!r} must return a Tensor.")
         if tuple(X.shape) != self.shape:
@@ -56,7 +59,7 @@ class LazyTensorReference:
             digest = tensor_sha256(X)
             if digest != self.sha256:
                 raise ValueError(f"Stored tensor {self.key!r} failed SHA-256 verification.")
-        return X  # self.shape
+        return X  # (...), equal to self.shape
 
 
 TensorValue = Tensor | LazyTensorReference
@@ -176,11 +179,75 @@ class EmbeddingBatch:
     attentions: Tensor | tuple[Tensor, ...] | None = None
 
 
+@dataclass(frozen=True, slots=True)
+class TapRecord:
+    """One sequence's outputs from a tap plan, keyed by tap name."""
+
+    id: str
+    sequence: str
+    tensors: Mapping[str, Tensor | TopKRow]
+    retained_positions: tuple[int, ...] | None = None
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.id, str) or not self.id:
+            raise ValueError("TapRecord.id must be a non-empty string.")
+        if not isinstance(self.sequence, str) or not self.sequence:
+            raise ValueError("TapRecord.sequence must be a non-empty string.")
+        if not isinstance(self.tensors, Mapping) or not self.tensors:
+            raise TypeError("TapRecord.tensors must be a non-empty mapping of tap name to Tensor.")
+        if not all(
+            isinstance(name, str) and isinstance(value, (Tensor, TopKRow))
+            for name, value in self.tensors.items()
+        ):
+            raise TypeError("TapRecord.tensors must map tap names to Tensor or TopKRow values.")
+        object.__setattr__(self, "tensors", MappingProxyType(dict(self.tensors)))
+        if self.retained_positions is not None:
+            positions = self.retained_positions
+            if (type(positions) is not tuple or not positions
+                    or any(type(p) is not int or not 0 <= p < len(self.sequence) for p in positions)
+                    or tuple(sorted(set(positions))) != positions):
+                raise ValueError(
+                    "TapRecord retained positions must be ordered original-sequence indices."
+                )
+
+
+@dataclass(frozen=True, slots=True)
+class TapRunReceipt:
+    """Completed sink delivery, retaining run metadata but no output tensors."""
+
+    record_count: int
+    metadata: Mapping[str, Any]
+
+
+class TapResult:
+    """Ordered tap records and the metadata needed to reproduce them."""
+
+    def __init__(
+        self,
+        records: Sequence[TapRecord],
+        metadata: Mapping[str, Any] | None = None,
+    ) -> None:
+        self.records: tuple[TapRecord, ...] = tuple(records)
+        self.metadata = dict(metadata or {})
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    def __iter__(self) -> Iterator[TapRecord]:
+        return iter(self.records)
+
+    def __getitem__(self, index: int) -> TapRecord:
+        return self.records[index]
+
+
 __all__ = [
     "EmbeddingBatch",
     "EmbeddingInput",
     "EmbeddingRecord",
     "EmbeddingResult",
     "LazyTensorReference",
+    "TapRecord",
+    "TapResult",
+    "TapRunReceipt",
     "TensorValue",
 ]

fastplms/features/__init__.py

@@ -0,0 +1,77 @@
+"""The FastPLMs feature store: one storage format for every embedding this workspace keeps.
+
+A feature is one value per sequence, addressed by the SHA-256 of the sequence and by a key that
+names the model, its revision, the sparse autoencoder, the layer, the pooling, the dtype, and the
+residue limit. `store` holds the directory format and `layouts` the row layouts: dense vectors,
+compressed-sparse pooled rows, ragged hidden states, and ragged top-k residue codes. `reader`
+serves rows by random access to loops that read a batch at a time, `writing` commits a window of
+rows as one segment, and `conversion` moves a cache another format holds into a store and proves the
+rows survived.
+"""
+
+from .async_writer import AsyncFeatureWriter, PackedBatch
+from .conversion import (
+    ConversionMismatch,
+    ConversionReceipt,
+    conversion_fingerprint,
+    convert_rows,
+    describe_file,
+)
+from .layouts import (
+    CSR,
+    DENSE,
+    LAYOUT_NAMES,
+    RAGGED,
+    RAGGED_TOPK,
+    SparseRow,
+    TopKRow,
+    dtype_name,
+    value_dtype,
+)
+from .reader import CsrRows, FeatureReader
+from .store import (
+    FORMAT,
+    FeatureStore,
+    RowAddress,
+    SegmentReceipt,
+    SegmentWriter,
+    StoredFeature,
+    features_in,
+    open_feature,
+    partition_sequences,
+    sequence_digest,
+)
+from .writing import write_rows
+
+
+__all__ = [
+    "CSR",
+    "DENSE",
+    "FORMAT",
+    "LAYOUT_NAMES",
+    "RAGGED",
+    "RAGGED_TOPK",
+    "AsyncFeatureWriter",
+    "ConversionMismatch",
+    "ConversionReceipt",
+    "CsrRows",
+    "FeatureReader",
+    "FeatureStore",
+    "PackedBatch",
+    "RowAddress",
+    "SegmentReceipt",
+    "SegmentWriter",
+    "SparseRow",
+    "StoredFeature",
+    "TopKRow",
+    "conversion_fingerprint",
+    "convert_rows",
+    "describe_file",
+    "dtype_name",
+    "features_in",
+    "open_feature",
+    "partition_sequences",
+    "sequence_digest",
+    "value_dtype",
+    "write_rows",
+]

fastplms/features/async_writer.py

@@ -0,0 +1,375 @@
+"""Write the rows of many streams behind the model: one bounded queue, one ingest thread, parallel part writes.
+
+The embedding loop runs on the device. This writer takes each finished batch from pinned host buffers,
+packs rows into parts of about ``part_bytes``, writes and hashes each part in one pass on a pool thread,
+and commits every stream together once ``segment_bytes`` of parts exist, so a killed run loses at most
+one segment. ``submit`` blocks only when the queued batches exceed ``queue_bytes``, which is what keeps
+the device from outrunning the disk.
+
+Symbols: b sequences of a batch; n token rows of a batch (sum of l_i + 2, with l_i the residues of
+sequence i after the crop); w stored columns of a stream; k SAE codes kept per token; c SAE codebook.
+Layouts of one batch, by stream:
+
+- ragged       values (n, w): row 0 CLS, rows 1..l_i residues, row l_i + 1 EOS of each sequence, in order.
+- ragged_topk  indices, values (n, k): the top-k SAE codes of every token row.
+- dense        values (b, w): one pooled vector per sequence.
+- csr          values (b, c) dense on arrival, compressed here to (nnz,) indices and values.
+"""
+
+from __future__ import annotations
+
+import threading
+import torch
+
+from collections import deque
+from collections.abc import Callable, Mapping, Sequence
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import ExitStack
+from dataclasses import dataclass, field
+from typing import Any
+from torch import Tensor
+
+from .layouts import CSR, DENSE, INDEX_DTYPE, OFFSET_DTYPE, RAGGED, RAGGED_TOPK
+from .store import FeatureStore, SegmentReceipt, SegmentWriter
+
+
+@dataclass(eq=False)
+class PackedBatch:
+    """One batch of finished rows for every stream, on the host, as the executor hands them over.
+
+    ``streams`` maps a stream to its host tensors by layout: ragged ``{"values": (n, w)}``, ragged
+    top-k ``{"values": (n, k), "indices": (n, k)}``, dense ``{"values": (b, w)}``, csr
+    ``{"values": (b, c)}`` dense, which the writer compresses. ``rows`` counts the stored rows of each
+    sequence in a ragged stream (l_i + 2 when the special tokens are kept). ``wait`` blocks until the
+    device copies landed and the batch passed its finite check, and raises if it did not.
+    """
+
+    sequences: tuple[str, ...]  # (b,) the exact text of each row, for the row identities
+    digests: tuple[str, ...]  # (b,) SHA-256 row keys, hashed once by the caller
+    rows: tuple[int, ...]  # (b,) stored rows per sequence for ragged streams
+    streams: Mapping[str, Mapping[str, Tensor]]  # per stream: values (n, w) | (n, k) with indices (n, k) | (b, w) | (b, c)
+    wait: Callable[[], None]
+    nbytes: int = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.nbytes = sum(
+            tensor.numel() * tensor.element_size() for group in self.streams.values() for tensor in group.values()
+        )
+
+
+@dataclass(eq=False)
+class _StreamBuffer:
+    """Rows of one stream waiting to fill a part."""
+
+    sequences: list[str] = field(default_factory=list)
+    digests: list[str] = field(default_factory=list)
+    rows: list[int] = field(default_factory=list)  # stored rows per sequence, as in PackedBatch
+    nnz: list[int] = field(default_factory=list)  # csr entries per sequence, zero for other layouts
+    tensors: dict[str, list[Tensor]] = field(default_factory=dict)  # per name: slices (n_i, w) or (b_i, w), joined at flush
+    nbytes: int = 0
+
+
+def _row_bytes(layout: str, tensors: Mapping[str, Tensor], rows: Sequence[int], nnz: Sequence[int]) -> list[int]:
+    """Encoded payload of each sequence's row, its offset included, which the part budget counts."""
+    # tensors: (b, w) per stream when dense, (nnz,) when csr, (n, w) when ragged; n = sum(rows) token rows.
+    if layout == DENSE:
+        per_row = sum(tensor.shape[1] * tensor.element_size() for tensor in tensors.values())
+        return [per_row] * len(rows)
+    if layout == CSR:
+        per_entry = sum(tensor.element_size() for tensor in tensors.values())
+        return [8 + count * per_entry for count in nnz]
+    per_token = sum(tensor.shape[1] * tensor.element_size() for tensor in tensors.values())
+    return [8 + count * per_token for count in rows]
+
+
+class AsyncFeatureWriter:
+    """Take finished batches from the embedding loop and commit them as rolling segments of every stream.
+
+    ``stores`` maps a stream to its opened store, whose descriptor carries the layout. ``wanted`` maps a
+    stream to the digests it lacks, so a rerun fills a lagging stream without rewriting the others.
+    ``row_records(stream, sequences, digests)`` returns each row's persisted identity. ``before_commit`` runs once
+    per committed group, before the first stream's marker, and may raise to refuse the commit.
+    """
+
+    def __init__(
+        self,
+        stores: Mapping[str, FeatureStore],
+        *,
+        fingerprint: str,
+        metadata: Mapping[str, Any],
+        wanted: Mapping[str, frozenset[str]],
+        row_records: Callable[[str, Sequence[str], Sequence[str]], Sequence[Mapping[str, Any]]],
+        part_bytes: int,
+        segment_bytes: int,
+        queue_bytes: int,
+        before_commit: Callable[[], None] | None = None,
+        workers: int = 4,
+        progress: Callable[[int], None] | None = None,
+        verify_staged: bool = False,
+    ) -> None:
+        if min(part_bytes, segment_bytes, queue_bytes, workers) < 1:
+            raise ValueError("part_bytes, segment_bytes, queue_bytes and workers must be positive.")
+        self._stores = dict(stores)
+        self._fingerprint = fingerprint
+        self._metadata = dict(metadata)
+        self._wanted = dict(wanted)
+        self._row_records = row_records
+        self._part_bytes = part_bytes
+        self._segment_bytes = segment_bytes
+        self._queue_bytes = queue_bytes
+        self._before_commit = before_commit
+        self._progress = progress
+        self._verify_staged = verify_staged
+        self._buffers = {name: _StreamBuffer() for name in self._stores}
+        self._pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="feature-part")
+        # Parts in flight, written or waiting for a pool thread; more would only queue host memory.
+        self._slots = threading.Semaphore(workers + 2)
+        self._futures: deque[Future[None]] = deque()
+        self._condition = threading.Condition()
+        self._queue: deque[PackedBatch | None] = deque()
+        self._queued_bytes = 0
+        self._error: BaseException | None = None
+        # Held around each commit, so an abort lands before a commit starts or after it ends, never during one.
+        self._commit_lock = threading.Lock()
+        self._aborted = False
+        self._segment_index = 0
+        self._segment_written = 0
+        self._receipts: dict[str, list[SegmentReceipt]] = {name: [] for name in self._stores}
+        self._stack = ExitStack()
+        self._writers: dict[str, SegmentWriter] = {}
+        self._checked = False
+        self._open_segments()
+        self._thread = threading.Thread(target=self._ingest, name="feature-ingest", daemon=True)
+        self._thread.start()
+
+    def submit(self, batch: PackedBatch) -> None:
+        """Queue one batch, blocking while the queue holds more than ``queue_bytes`` of finished rows."""
+        with self._condition:
+            while self._error is None and self._queued_bytes and self._queued_bytes + batch.nbytes > self._queue_bytes:
+                self._condition.wait()
+            self._raise_error()
+            self._queue.append(batch)
+            self._queued_bytes += batch.nbytes
+            self._condition.notify_all()
+
+    def close(self) -> dict[str, list[SegmentReceipt]]:
+        """Write what is queued, commit the last segment, and return every committed segment by stream."""
+        with self._condition:
+            self._queue.append(None)
+            self._condition.notify_all()
+        self._thread.join()
+        self._pool.shutdown(wait=True)
+        if self._error is not None:
+            self._discard()
+            raise self._error
+        return self._receipts
+
+    def abort(self) -> None:
+        """Stop without committing: queued rows are dropped and open segments stay uncommitted for ``sweep``."""
+        with self._commit_lock:  # a commit in flight finishes; none starts after this
+            self._aborted = True
+        with self._condition:
+            self._error = self._error or RuntimeError("The feature writer was aborted.")
+            self._queue.clear()
+            self._condition.notify_all()
+        self._thread.join()
+        self._pool.shutdown(wait=True)
+        self._discard()
+
+    def _discard(self) -> None:
+        """Close the open segments without committing them; a closed writer is skipped by its context."""
+        with self._commit_lock:  # one closer at a time, and never during a commit
+            for writer in self._writers.values():
+                writer.closed = True
+            self._stack.close()
+
+    def _raise_error(self) -> None:
+        if self._error is not None:
+            raise self._error
+
+    def _open_segments(self) -> None:
+        name = f"{self._fingerprint}-{self._segment_index:05d}"
+        for stream, store in self._stores.items():
+            self._writers[stream] = self._stack.enter_context(store.segment(
+                name, {**self._metadata, "segment_index": self._segment_index},
+                before_commit=self._check_once, verify_staged=self._verify_staged,
+            ))
+
+    def _check_once(self) -> None:
+        """Run the caller's pre-commit check for the first stream of a group; its siblings reuse the verdict."""
+        if not self._checked:
+            if self._before_commit is not None:
+                self._before_commit()
+            self._checked = True
+
+    def _ingest(self) -> None:
+        try:
+            while True:
+                with self._condition:
+                    while not self._queue and self._error is None:
+                        self._condition.wait()
+                    if self._error is not None:
+                        return
+                    batch = self._queue.popleft()
+                if batch is None:
+                    self._commit_group()
+                    return
+                batch.wait()  # the device copies landed and the finite check passed
+                for stream in self._stores:
+                    self._add(stream, batch)
+                with self._condition:
+                    self._queued_bytes -= batch.nbytes
+                    self._condition.notify_all()
+                if self._progress is not None:
+                    self._progress(len(batch.sequences))
+                if self._segment_written >= self._segment_bytes:
+                    self._commit_group(reopen=True)
+        # A worker thread has no caller to raise to: keep the error for `_raise_error` on the caller's thread.
+        except BaseException as error:  # noqa: broad-except
+            with self._condition:
+                self._error = self._error or error
+                self._condition.notify_all()
+
+    def _add(self, stream: str, batch: PackedBatch) -> None:
+        layout = self._stores[stream].spec.layout
+        keep = [index for index, digest in enumerate(batch.digests) if digest in self._wanted[stream]]
+        if not keep:
+            return
+        tensors = batch.streams[stream]  # values (n, w) | (n, k) with indices (n, k) | (b, w) | (b, c), by layout
+        rows = list(batch.rows)  # (b,) stored rows per sequence: l_i + 2 when CLS and EOS are kept
+        if layout == CSR:
+            tensors, nnz = _compress_csr(tensors["values"])  # indices, values (nnz,); nnz per sequence (b,)
+        else:
+            nnz = [0] * len(rows)
+        if len(keep) != len(rows):
+            tensors, rows, nnz = _select_rows(layout, tensors, rows, nnz, keep)
+        sequences = [batch.sequences[index] for index in keep]  # (m,) m sequences this stream lacks
+        digests = [batch.digests[index] for index in keep]  # (m,)
+        sizes = _row_bytes(layout, tensors, rows, nnz)  # (m,) payload bytes per sequence
+        if any(size > self._part_bytes for size in sizes):
+            raise ValueError("A feature row exceeds max_part_bytes; increase the explicit part budget.")
+        buffer = self._buffers[stream]
+        start = 0
+        while start < len(sequences):
+            # Take rows until the next would overflow the part, flush, and go on; a flushed buffer fits any row.
+            room = self._part_bytes - buffer.nbytes
+            stop, used = start, 0
+            while stop < len(sequences) and used + sizes[stop] <= room:
+                used += sizes[stop]
+                stop += 1
+            if stop == start:
+                self._flush(stream)
+                buffer = self._buffers[stream]
+                continue
+            _take(buffer, layout, tensors, sequences, digests, rows, nnz, start, stop, used)
+            self._segment_written += used  # buffered rows count toward the segment, so a small segment commits per batch
+            if buffer.nbytes >= self._part_bytes:
+                self._flush(stream)
+                buffer = self._buffers[stream]
+            start = stop
+
+    def _flush(self, stream: str) -> None:
+        buffer = self._buffers[stream]
+        if not buffer.sequences:
+            return
+        self._buffers[stream] = _StreamBuffer()
+        # A failed part surfaces at the next flush, not only at the end of the segment.
+        while self._futures and self._futures[0].done():
+            self._futures.popleft().result()
+        writer = self._writers[stream]
+        part = writer.reserve_part()
+        self._slots.acquire()
+        future = self._pool.submit(self._write_part, stream, writer, part, buffer)
+        future.add_done_callback(lambda _: self._slots.release())
+        self._futures.append(future)
+
+    def _write_part(self, stream: str, writer: SegmentWriter, part: int, buffer: _StreamBuffer) -> None:
+        spec = self._stores[stream].spec
+        tensors = _pack(spec.layout, buffer)  # offsets (b + 1,) and values (n, w), or (b, w), or indptr (b + 1,) and (nnz,)
+        identities = self._row_records(stream, buffer.sequences, buffer.digests)  # (b,) one identity per sequence
+        residues = buffer.rows if spec.layout in (RAGGED, RAGGED_TOPK) else [0] * len(buffer.sequences)  # (b,) stored rows
+        writer.append_packed(part, buffer.digests, tensors, residues, row_metadata=identities)
+
+    def _drain(self) -> None:
+        """Wait for every part in flight, surfacing the first write error."""
+        while self._futures:
+            self._futures.popleft().result()
+
+    def _commit_group(self, *, reopen: bool = False) -> None:
+        """Flush every stream, commit its open segment, and optionally open the next segment group."""
+        for stream in self._stores:
+            self._flush(stream)
+        self._drain()
+        self._checked = False
+        for stream, writer in self._writers.items():
+            with self._commit_lock:
+                if self._aborted:
+                    raise RuntimeError("The feature writer was aborted before its segment committed.")
+                if writer.parts:
+                    self._receipts[stream].append(writer.commit())
+                else:
+                    writer.abandon()
+        self._stack.close()
+        if reopen:
+            self._stack = ExitStack()
+            self._segment_index += 1
+            self._segment_written = 0
+            self._open_segments()
+
+
+def _compress_csr(dense: Tensor) -> tuple[dict[str, Tensor], list[int]]:
+    """Keep the nonzero codes of each sequence's row, in row-major order, as indices and values."""
+    # dense: (b, c) float32, zero where no code fired anywhere in the sequence
+    present = dense != 0  # (b, c)
+    counts = present.sum(dim=1).tolist()  # (b,) nnz per sequence
+    columns = present.nonzero()[:, 1].to(INDEX_DTYPE)  # (nnz,)
+    return {"indices": columns, "values": dense[present]}, [int(count) for count in counts]  # values (nnz,)
+
+
+def _select_rows(
+    layout: str, tensors: Mapping[str, Tensor], rows: Sequence[int], nnz: Sequence[int], keep: Sequence[int],
+) -> tuple[dict[str, Tensor], list[int], list[int]]:
+    """Gather the sequences at ``keep`` from packed batch tensors, for a rerun that fills only some rows."""
+    # tensors: (b, w) when dense, (n, w) when ragged, (nnz,) when csr; the kept rows keep the layout.
+    if layout == DENSE:
+        index = torch.tensor(list(keep), dtype=torch.int64)  # (m,) m kept sequences
+        picked = {name: tensor[index] for name, tensor in tensors.items()}
+        return picked, [rows[i] for i in keep], [0] * len(keep)  # (m, w) per stream, m kept sequences
+    spans = rows if layout in (RAGGED, RAGGED_TOPK) else nnz
+    pieces = {name: torch.split(tensor, list(spans)) for name, tensor in tensors.items()}  # (span_i, ...) per sequence
+    selected = {name: torch.cat([parts[i] for i in keep]) for name, parts in pieces.items()}
+    return selected, [rows[i] for i in keep], [nnz[i] for i in keep]  # (n_kept, w) or (nnz_kept,) per stream
+
+
+def _take(
+    buffer: _StreamBuffer, layout: str, tensors: Mapping[str, Tensor], sequences: Sequence[str],
+    digests: Sequence[str], rows: Sequence[int], nnz: Sequence[int], start: int, stop: int, used: int,
+) -> None:
+    """Move sequences ``[start, stop)`` of a batch into the stream's part buffer, slicing the packed tensors."""
+    # tensors: (b, w) when dense, (n, w) when ragged, (nnz,) when csr; rows low:high of each go to the buffer.
+    buffer.sequences.extend(sequences[start:stop])
+    buffer.digests.extend(digests[start:stop])
+    buffer.rows.extend(rows[start:stop])
+    buffer.nnz.extend(nnz[start:stop])
+    if layout == DENSE:
+        low, high = start, stop  # one row per sequence
+    else:
+        spans = rows if layout in (RAGGED, RAGGED_TOPK) else nnz
+        low, high = sum(spans[:start]), sum(spans[:stop])  # first and last packed row of the slice
+    for name, tensor in tensors.items():
+        buffer.tensors.setdefault(name, []).append(tensor[low:high])
+    buffer.nbytes += used
+
+
+def _pack(layout: str, buffer: _StreamBuffer) -> dict[str, Tensor]:
+    """Concatenate a part's batches into the layout's tensors, with offsets naming each sequence's span."""
+    joined = {name: torch.cat(parts) for name, parts in buffer.tensors.items()}  # (n, ...) one copy per tensor
+    if layout == DENSE:
+        return joined  # values (b, w)
+    spans = torch.tensor(buffer.rows if layout in (RAGGED, RAGGED_TOPK) else buffer.nnz, dtype=OFFSET_DTYPE)  # (b,)
+    offsets = torch.zeros(len(spans) + 1, dtype=OFFSET_DTYPE)  # (b + 1,)
+    offsets[1:] = torch.cumsum(spans, dim=0)
+    return {("indptr" if layout == CSR else "offsets"): offsets, **joined}  # offsets (b + 1,); values (n, w) or (nnz,)
+
+
+__all__ = ["AsyncFeatureWriter", "PackedBatch"]

fastplms/features/conversion.py

@@ -0,0 +1,272 @@
+"""Checked conversion of a cache another format holds into a feature store.
+
+A project that already holds embeddings in its own format keeps them by converting them, not by
+re-embedding. ``convert_rows`` takes the old cache as a function that decodes it, writes the rows
+through the store's ordinary segment writer, then decodes the old cache a second time and compares
+every row it produced with what the store now returns. Comparison is bit-exact on the stored
+representation, so a converter never has to argue that two floats are close enough.
+
+Nothing is repaired or skipped quietly. A row the store's dtype cannot hold exactly, a sequence
+absent after the commit, a row that differs, a source that repeats a sequence with different rows,
+and a source that changes between its two decodes each raise ``ConversionMismatch``. The old cache is
+never modified or deleted: the move is undone by removing the segments this call names, and the
+origin recorded in each commit marker says which files those rows came from.
+
+What a project supplies is only the decoder of its old format. The decoder yields
+``(sequence, row)`` pairs, where a row is a tensor for ``dense`` and ``ragged`` features, a
+``SparseRow`` for ``csr``, and a ``TopKRow`` for ``ragged_topk``. It takes no arguments, because it
+is called twice.
+"""
+
+from __future__ import annotations
+
+import json
+import torch
+
+from collections.abc import Callable, Iterable, Iterator, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from torch import Tensor
+
+from .digests import file_sha256, json_sha256
+from .layouts import CSR, RAGGED_TOPK, SparseRow, TopKRow
+from .reader import FeatureReader
+from .store import (
+    COMMIT_FILE,
+    SEGMENTS_DIRECTORY,
+    FeatureStore,
+    StoredFeature,
+    sequence_digest,
+)
+
+
+Row = Tensor | SparseRow | TopKRow
+
+
+class ConversionMismatch(ValueError):
+    """The store does not hold exactly what the old cache held."""
+
+
+@dataclass(frozen=True, slots=True)
+class ConversionReceipt:
+    """What one call did: the segments it committed and how many rows it compared."""
+
+    segments: tuple[str, ...]
+    source_rows: int
+    written: int
+    skipped: int
+    verified: int
+
+
+def describe_file(path: str | Path) -> dict[str, Any]:
+    """A file as an origin records it: its name, size, and content digest."""
+
+    location = Path(path)
+    return {
+        "name": location.name,
+        "bytes": location.stat().st_size,
+        "sha256": file_sha256(location),
+    }
+
+
+def conversion_fingerprint(origin: Mapping[str, Any]) -> str:
+    """The stable name of the conversion of this origin, from which segment names derive."""
+
+    return "convert-" + json_sha256(dict(origin), allow_nan=False)[:16]
+
+
+def convert_rows(
+    store: FeatureStore,
+    source: Callable[[], Iterable[tuple[str, Row]]],
+    *,
+    origin: Mapping[str, Any],
+    rows_per_segment: int = 100_000,
+    window_rows: int = 1_024,
+    max_tensor_bytes: int = 256 * 1024**2,
+) -> ConversionReceipt:
+    """Write the decoded rows into ``store``, then prove the store holds them.
+
+    ``origin`` is plain data naming where the rows came from, normally the old cache's format, the
+    decoder, and ``describe_file`` of each file. It is recorded in every segment's commit marker and
+    names the conversion, so calling this again with the same origin resumes an interrupted
+    conversion: committed segments are kept, sequences already in the store are not written twice,
+    and the comparison runs over the whole source. A long conversion commits a segment once it holds
+    ``rows_per_segment`` new rows, rounded up to a window, so an interruption loses at most one
+    segment.
+
+    Rows the store held before the call are compared too. A store that mixes converted rows with
+    rows embedded afresh therefore fails here, which is the point: the old numbers are not the
+    store's numbers.
+    """
+
+    if rows_per_segment < 1 or window_rows < 1:
+        raise ValueError("rows_per_segment and window_rows must be positive.")
+    fingerprint = conversion_fingerprint(origin)
+    ordinal = _committed_segments(store, fingerprint)
+    segments: list[str] = []
+    source_rows = 0
+    written = 0
+
+    def counted() -> Iterator[tuple[str, Row]]:
+        nonlocal source_rows
+        for pair in source():
+            source_rows += 1
+            yield pair
+
+    windows = _unique_windows(store.spec, counted(), window_rows)
+    pending = next(windows, None)
+    while pending is not None:
+        name = f"{fingerprint}-{ordinal:05d}"
+        metadata = {"conversion": {"origin": dict(origin), "ordinal": ordinal}}
+        staged: set[str] = set()  # digests written into this segment, which is not yet committed
+        with store.segment(name, metadata) as writer:
+            while pending is not None and len(staged) < rows_per_segment:
+                absent = set(store.missing([sequence for sequence, _ in pending]))
+                fresh = [
+                    pair for pair in pending
+                    if pair[0] in absent and sequence_digest(pair[0]) not in staged
+                ]
+                if fresh:
+                    writer.append_bounded(
+                        [sequence for sequence, _ in fresh], [row for _, row in fresh],
+                        max_tensor_bytes=max_tensor_bytes,
+                    )
+                    staged.update(sequence_digest(sequence) for sequence, _ in fresh)
+                pending = next(windows, None)
+        if staged:
+            segments.append(name)
+            ordinal += 1
+            written += len(staged)
+
+    if source_rows == 0:
+        raise ConversionMismatch("The source decoded no rows; check the path and the decoder.")
+    verified = _verify(store, source, window_rows)
+    if verified != source_rows:
+        raise ConversionMismatch(
+            f"The source decoded {source_rows} rows the first time and {verified} the second."
+        )
+    return ConversionReceipt(
+        segments=tuple(segments), source_rows=source_rows, written=written,
+        skipped=source_rows - written, verified=verified,
+    )
+
+
+def _committed_segments(store: FeatureStore, fingerprint: str) -> int:
+    directory = store.directory / SEGMENTS_DIRECTORY
+    return sum(
+        1 for marker in directory.glob(f"{fingerprint}-*/{COMMIT_FILE}") if marker.is_file()
+    )
+
+
+def _windows(rows: Iterable[tuple[str, Row]], size: int) -> Iterator[list[tuple[str, Row]]]:
+    window: list[tuple[str, Row]] = []
+    for pair in rows:
+        window.append(pair)
+        if len(window) == size:
+            yield window
+            window = []
+    if window:
+        yield window
+
+
+def _unique_windows(
+    spec: StoredFeature, rows: Iterable[tuple[str, Row]], size: int,
+) -> Iterator[list[tuple[str, Row]]]:
+    """Windows in which each sequence appears once.
+
+    A sequence repeated inside a window keeps its first row, and is an error if the rows differ.
+    A repeat in a later window is skipped when its sequence is already written, and the comparison
+    after the commit raises if its row differs from the one that was.
+    """
+
+    for window in _windows(rows, size):
+        unique: dict[str, tuple[str, Row]] = {}
+        for sequence, row in window:
+            digest = sequence_digest(sequence)
+            first = unique.setdefault(digest, (sequence, row))
+            if first[1] is not row and _row_bytes(first[1], spec) != _row_bytes(row, spec):
+                raise ConversionMismatch(
+                    f"The source repeats a sequence (sha256 {digest[:12]}) with different rows."
+                )
+        yield list(unique.values())
+
+
+def _verify(
+    store: FeatureStore, source: Callable[[], Iterable[tuple[str, Row]]], window_rows: int,
+) -> int:
+    """Decode the source again and compare each row with a fresh read of the committed store."""
+
+    spec = store.spec
+    verified = 0
+    with FeatureReader.open(store.directory) as reader:
+        for window in _windows(source(), window_rows):
+            sequences = [sequence for sequence, _ in window]
+            absent = reader.missing(sequences)
+            if absent:
+                raise ConversionMismatch(
+                    f"{len(absent)} converted sequences are absent from {spec.key!r} after commit."
+                )
+            for (sequence, row), stored in zip(window, _read(reader, sequences), strict=True):
+                if _row_bytes(row, spec) != _row_bytes(stored, spec):
+                    raise ConversionMismatch(
+                        f"Feature {spec.key!r} differs from the source for the sequence with "
+                        f"sha256 {sequence_digest(sequence)[:12]}."
+                    )
+            verified += len(window)
+    return verified
+
+
+def _read(reader: FeatureReader, sequences: list[str]) -> list[Row]:
+    layout = reader.spec.layout
+    if layout == CSR:
+        return list(reader.read_sparse(sequences))
+    if layout == RAGGED_TOPK:
+        return list(reader.read_topk(sequences))
+    return list(reader.read(sequences))
+
+
+def _row_bytes(row: Row, spec: StoredFeature) -> bytes:
+    """A row as the store holds it, as bytes: shapes, then each tensor's raw storage.
+
+    Values are cast to the feature's dtype first, and a cast that changes any value raises, so
+    equal bytes mean equal stored rows.
+    """
+
+    # row: (w,) dense, (r_i, d) ragged; a SparseRow holds (nnz,) tensors and a TopKRow (r_i, k) tensors
+    if isinstance(row, SparseRow):
+        tensors = [row.indices.to(torch.int32), _stored_values(row.values, spec.dtype)]  # (nnz,), (nnz,)
+        if row.positions is not None:
+            tensors.append(row.positions.to(torch.int16))  # (nnz,)
+    elif isinstance(row, TopKRow):
+        tensors = [row.indices.to(torch.int32), _stored_values(row.values, spec.dtype)]  # (r_i, k), (r_i, k)
+    else:
+        tensors = [_stored_values(row, spec.dtype)]  # (w,) or (r_i, d)
+    shapes = json.dumps([list(tensor.shape) for tensor in tensors]).encode("utf-8")
+    return shapes + b"".join(_raw_bytes(tensor) for tensor in tensors)
+
+
+def _stored_values(values: Tensor, dtype: torch.dtype) -> Tensor:
+    # values: (...), the shape of one stored tensor; the cast keeps it
+    cast = values.detach().to("cpu").to(dtype)  # (...)
+    if values.dtype != dtype and _raw_bytes(cast.to(values.dtype)) != _raw_bytes(values.detach().cpu()):
+        raise ConversionMismatch(
+            f"Storing {values.dtype} values as {dtype} would change them; choose a wider dtype."
+        )
+    return cast  # (...)
+
+
+def _raw_bytes(tensor: Tensor) -> bytes:
+    # tensor: (...)
+    flat = tensor.detach().to("cpu").contiguous().reshape(-1)  # (n,), n = tensor.numel()
+    # An empty tensor can carry a zero stride, which the byte view refuses, and has no bytes to compare.
+    return b"" if flat.numel() == 0 else flat.view(torch.uint8).numpy().tobytes()
+
+
+__all__ = [
+    "ConversionMismatch",
+    "ConversionReceipt",
+    "convert_rows",
+    "conversion_fingerprint",
+    "describe_file",
+]

fastplms/features/digests.py

@@ -0,0 +1,35 @@
+"""SHA-256 digests of files and JSON values, the identities FastPLMs records and compares.
+
+This file exists twice, byte for byte: here and as ``features/digests.py``. ``features`` loads as a standalone
+package (``foundry.embedding.private_store``) and imports nothing outside itself, so it cannot reach this
+module. ``tests/tier1_unit/test_features_package_is_self_contained.py`` fails when the two differ.
+"""
+
+from __future__ import annotations
+
+import hashlib
+
+from pathlib import Path
+from typing import Any
+
+from .json_files import compact_json
+
+
+FILE_READ_BYTES = 1024 * 1024
+
+
+def file_sha256(path: str | Path) -> str:
+    """Return the SHA-256 of a file's bytes, read in blocks so a checkpoint never sits in memory."""
+
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as handle:
+        while block := handle.read(FILE_READ_BYTES):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def json_sha256(value: Any, *, ensure_ascii: bool = True, allow_nan: bool = True) -> str:
+    """Return the SHA-256 of ``value`` in its compact, key-sorted JSON form (``compact_json``)."""
+
+    encoded = compact_json(value, ensure_ascii=ensure_ascii, allow_nan=allow_nan).encode("utf-8")
+    return hashlib.sha256(encoded).hexdigest()

fastplms/features/json_files.py

@@ -0,0 +1,45 @@
+"""The two JSON text forms FastPLMs writes: compact for hashing, indented for files people read.
+
+This file exists twice, byte for byte: here and as ``features/json_files.py``. ``features`` loads as a
+standalone package (``foundry.embedding.private_store``) and imports nothing outside itself, so it cannot
+reach this module. ``tests/tier1_unit/test_features_package_is_self_contained.py`` fails when the two differ.
+"""
+
+from __future__ import annotations
+
+import json
+
+from typing import Any
+
+
+def compact_json(value: Any, *, ensure_ascii: bool = True, allow_nan: bool = True) -> str:
+    """Serialize with sorted keys and no whitespace, the form a digest or an identity is taken over."""
+
+    return json.dumps(
+        value,
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=ensure_ascii,
+        allow_nan=allow_nan,
+    )
+
+
+def indented_json(
+    value: Any,
+    *,
+    ensure_ascii: bool = True,
+    allow_nan: bool = True,
+    sort_keys: bool = True,
+) -> str:
+    """Serialize with two-space indentation and one trailing newline, the form of a stored JSON file."""
+
+    return (
+        json.dumps(
+            value,
+            indent=2,
+            sort_keys=sort_keys,
+            ensure_ascii=ensure_ascii,
+            allow_nan=allow_nan,
+        )
+        + "\n"
+    )

fastplms/features/layouts.py

@@ -0,0 +1,354 @@
+"""The row layouts a feature segment stores, as memory-mappable safetensors tensors.
+
+A feature is one value per sequence, and the layouts differ in what that value is:
+
+- ``dense``: a fixed-width vector, ``values`` of shape ``(n, w)``. Pooled embeddings.
+- ``csr``: a sparse fixed-width vector, compressed by row. Max-pooled sparse-autoencoder codes,
+  where a row holds at most one entry per code that fired anywhere in the sequence. ``positions``
+  carries the argmax residue of each entry, which is what makes a code interpretable as "here",
+  and is optional because a pooling that discards it has nothing to store.
+- ``ragged``: per-row values, ``values`` of shape ``(sum r_i, d)`` with ``offsets`` naming each
+  sequence's span. Hidden states, where ``r_i`` is the sequence's stored row count: its biological residue
+  count ``l`` in a residue-only (``feature_spec_v1``) store, and ``l + 2`` in a canonical store, whose rows are
+  CLS, the ``l`` residues, then EOS.
+- ``ragged_topk``: per-row sparse codes, ``indices`` and ``values`` of shape
+  ``(sum r_i, k)`` with sequence ``offsets``, ``r_i`` as above. All k entries, including zeros, retain their
+  order. Unlike pooled csr positions, these rows retain every stored row and support crop-local pooling.
+
+Rows are addressed individually and never sliced by column, so compressed-sparse-row is the right
+sparse layout: a batch materializes with one gather and no search. The index dtypes follow
+enzyme_loop's measured store: int32 code indices, int16 residue positions, and int64 row offsets,
+which keep an entry to eight bytes against four bytes per column dense.
+
+Every layout stores its values in the feature's declared dtype, so a segment is a lossless record
+of what the model produced at that precision, and no reader has to guess.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from torch import Tensor
+
+
+DENSE = "dense"
+CSR = "csr"
+RAGGED = "ragged"
+RAGGED_TOPK = "ragged_topk"
+LAYOUT_NAMES = (DENSE, CSR, RAGGED, RAGGED_TOPK)
+
+INDEX_DTYPE = torch.int32
+POSITION_DTYPE = torch.int16
+OFFSET_DTYPE = torch.int64
+
+VALUE_DTYPES = (torch.float64, torch.float32, torch.float16, torch.bfloat16)
+DTYPE_NAMES = {dtype: str(dtype).removeprefix("torch.") for dtype in VALUE_DTYPES}
+NAME_DTYPES = {name: dtype for dtype, name in DTYPE_NAMES.items()}
+
+
+@dataclass(frozen=True, slots=True)
+class TopKRow:
+    """One sequence's sparse residue codes: integer indices and values, both ``(r_i, k)``.
+
+    The store validates the declared codebook and k. Zero-valued entries and code order are
+    retained exactly; there is deliberately no implicit residue-by-codebook densification.
+    """
+
+    indices: Tensor
+    values: Tensor
+
+
+@dataclass(frozen=True, slots=True)
+class SparseRow:
+    """One compressed row: the codes that fired, their values, and where each peaked.
+
+    ``indices`` and ``values`` have shape ``(nnz_i,)``. ``positions`` has the same shape when the
+    segment stores them, and is None when it does not.
+    """
+
+    indices: Tensor
+    values: Tensor
+    positions: Tensor | None
+
+    def to_dense(self, width: int) -> Tensor:
+        """The row as a ``(width,)`` vector, zero where no code fired."""
+
+        # self.indices, self.values: (nnz,), the entries that fired
+        dense = torch.zeros(width, dtype=self.values.dtype)  # (w,)
+        dense[self.indices.to(torch.int64)] = self.values
+        return dense  # (w,)
+
+    @classmethod
+    def from_dense(cls, vector: Tensor, positions: Tensor | None = None) -> SparseRow:
+        """Compress a ``(w,)`` vector by keeping its exactly non-zero entries.
+
+        Top-k sparse-autoencoder pooling leaves exact zeros where a code never fired, so this
+        drops nothing a reader could want. It is not a threshold: a code that fired weakly is
+        kept. ``positions`` is the full ``(w,)`` argmax residue vector, gathered at the same
+        entries.
+        """
+
+        # vector: (w,); positions: (w,) or None
+        if vector.ndim != 1:
+            raise ValueError(f"A sparse row comes from a one-dimensional vector; received {tuple(vector.shape)}.")
+        kept = torch.nonzero(vector, as_tuple=False).flatten()  # (nnz,)
+        if positions is not None and positions.shape != vector.shape:
+            raise ValueError(
+                f"positions must have the vector's shape {tuple(vector.shape)}; received "
+                f"{tuple(positions.shape)}."
+            )
+        return cls(
+            indices=kept.to(INDEX_DTYPE),  # (nnz,)
+            values=vector[kept],  # (nnz,)
+            positions=None if positions is None else positions[kept],  # (nnz,)
+        )
+
+
+def dtype_name(dtype: torch.dtype) -> str:
+    """The stored name of a value dtype, rejecting one no layout stores."""
+
+    if dtype not in DTYPE_NAMES:
+        raise ValueError(
+            f"A feature stores float64, float32, float16, or bfloat16 values; received {dtype}."
+        )
+    return DTYPE_NAMES[dtype]
+
+
+def value_dtype(name: str) -> torch.dtype:
+    """The dtype a stored name means."""
+
+    if name not in NAME_DTYPES:
+        raise ValueError(f"Unknown feature value dtype {name!r}; expected one of {list(NAME_DTYPES)}.")
+    return NAME_DTYPES[name]
+
+
+def encode_dense(rows: Sequence[Tensor], width: int, dtype: torch.dtype) -> dict[str, Tensor]:
+    """Stack ``(w,)`` rows into one ``values`` tensor of shape ``(n, w)``."""
+
+    # rows: (w,) each, n tensors; w = width
+    stacked = torch.empty((len(rows), width), dtype=dtype)  # (n, w)
+    for position, row in enumerate(rows):
+        vector = row.detach().to("cpu")  # (w,)
+        if vector.ndim != 1 or vector.shape[0] != width:
+            raise ValueError(
+                f"A dense feature row must have shape ({width},); row {position} has "
+                f"{tuple(vector.shape)}."
+            )
+        stacked[position] = vector.to(dtype)
+    return {"values": stacked}  # {"values": (n, w)}
+
+
+def row_tensor_bytes(
+    row: Tensor | SparseRow | TopKRow, layout: str, width: int, dtype: torch.dtype,
+    *, positions: bool,
+) -> int:
+    """Exact per-row encoded payload, including its offset but not the part's initial offset."""
+    # row: (w,) dense, (r_i, d) ragged, or a SparseRow or TopKRow
+    value_size = torch.empty(0, dtype=dtype).element_size()
+    if layout == RAGGED_TOPK:
+        if not isinstance(row, TopKRow):
+            raise TypeError("Ragged top-k payload sizing requires a TopKRow.")
+        return 8 + row.values.numel() * (4 + value_size)
+    if layout == CSR:
+        if not isinstance(row, SparseRow):
+            raise TypeError("CSR payload sizing requires a SparseRow.")
+        return 8 + row.values.numel() * (4 + value_size + (2 if positions else 0))
+    if not isinstance(row, Tensor):
+        raise TypeError("Dense and ragged payload sizing requires tensor rows.")
+    return width * value_size if layout == DENSE else 8 + row.numel() * value_size
+
+
+def encode_csr(
+    rows: Sequence[SparseRow], width: int, dtype: torch.dtype
+) -> dict[str, Tensor]:
+    """Concatenate sparse rows, with ``indptr`` naming each row's span.
+
+    Every row must agree about positions: either all carry them or none does, because one segment
+    stores one tensor set.
+    """
+
+    # rows[i].indices, .values, .positions: (nnz_i,); the segment holds n rows and nnz = sum(nnz_i) entries
+    with_positions = [row.positions is not None for row in rows]
+    if any(with_positions) and not all(with_positions):
+        raise ValueError(
+            "Either every sparse row carries argmax positions or none does; this batch mixes both."
+        )
+    indptr = torch.zeros(len(rows) + 1, dtype=OFFSET_DTYPE)  # (n + 1,)
+    indices: list[Tensor] = []
+    values: list[Tensor] = []
+    positions: list[Tensor] = []
+    for position, row in enumerate(rows):
+        if row.indices.dtype not in (torch.int16, torch.int32, torch.int64):
+            raise ValueError("Sparse code indices must be signed integer tensors.")
+        row_indices = row.indices.detach().to("cpu").to(torch.int64)  # (nnz_i,)
+        row_values = row.values.detach().to("cpu")  # (nnz_i,)
+        if row_indices.ndim != 1 or row_values.shape != row_indices.shape:
+            raise ValueError(
+                f"Sparse row {position} needs one-dimensional indices and values of equal length; "
+                f"received {tuple(row_indices.shape)} and {tuple(row_values.shape)}."
+            )
+        if row_indices.numel() and (int(row_indices.min()) < 0 or int(row_indices.max()) >= width):
+            raise ValueError(
+                f"Sparse row {position} names a code outside 0..{width - 1}."
+            )
+        if row_indices.numel() and int(row_indices.max()) > torch.iinfo(INDEX_DTYPE).max:
+            raise ValueError("Sparse code index exceeds the stored integer range.")
+        if row_indices.unique().numel() != row_indices.numel():
+            raise ValueError("Sparse code indices must be unique within each row.")
+        indptr[position + 1] = int(indptr[position]) + row_indices.numel()
+        indices.append(row_indices.to(INDEX_DTYPE))
+        values.append(row_values.to(dtype))
+        if row.positions is not None:
+            if row.positions.dtype not in (torch.int16, torch.int32, torch.int64):
+                raise ValueError("Sparse residue positions must be signed integer tensors.")
+            row_positions = row.positions.detach().to("cpu")  # (nnz_i,)
+            if row_positions.shape != row_indices.shape:
+                raise ValueError(
+                    f"Sparse row {position} has {row_positions.numel()} positions for "
+                    f"{row_indices.numel()} entries."
+                )
+            if row_positions.numel() and (
+                int(row_positions.min()) < 0
+                or int(row_positions.max()) > torch.iinfo(POSITION_DTYPE).max
+            ):
+                raise ValueError("Sparse residue position exceeds the stored integer range.")
+            positions.append(row_positions.to(POSITION_DTYPE))
+
+    tensors = {
+        "indptr": indptr,  # (n + 1,)
+        "indices": _concatenate(indices, INDEX_DTYPE),  # (nnz,)
+        "values": _concatenate(values, dtype),  # (nnz,)
+    }
+    if positions:
+        tensors["positions"] = _concatenate(positions, POSITION_DTYPE)  # (nnz,)
+    return tensors  # indptr (n + 1,); indices, values and positions (nnz,)
+
+
+def encode_ragged(rows: Sequence[Tensor], width: int, dtype: torch.dtype) -> dict[str, Tensor]:
+    """Concatenate ``(r_i, d)`` residue blocks, with ``offsets`` naming each sequence's span."""
+
+    # rows: (r_i, d) each, n tensors; d = width
+    offsets = torch.zeros(len(rows) + 1, dtype=OFFSET_DTYPE)  # (n + 1,)
+    blocks: list[Tensor] = []
+    for position, row in enumerate(rows):
+        block = row.detach().to("cpu")  # (r_i, d)
+        if block.ndim != 2 or block.shape[1] != width:
+            raise ValueError(
+                f"A ragged feature row must have shape (r_i, {width}); row {position} has "
+                f"{tuple(block.shape)}."
+            )
+        offsets[position + 1] = int(offsets[position]) + block.shape[0]
+        blocks.append(block.to(dtype))
+    values = (  # (sum r_i, d)
+        torch.cat(blocks, dim=0) if blocks else torch.empty((0, width), dtype=dtype)
+    )
+    return {"offsets": offsets, "values": values}  # {"offsets": (n + 1,), "values": (sum r_i, d)}
+
+
+def validate_topk(indices: Tensor, values: Tensor, width: int, count: int) -> None:
+    """Reject malformed residue codes before writing and when reading committed tensors."""
+    # indices, values: (r, k), with k = count; r residues of one sequence, or of every sequence of a part
+    if indices.dtype not in (torch.int16, torch.int32, torch.int64):
+        raise ValueError("Top-k code indices must be signed integer tensors.")
+    if (indices.ndim != 2 or indices.shape[1] != count or values.shape != indices.shape):
+        raise ValueError(f"Top-k indices and values must both have shape (residues, {count}).")
+    if not values.is_floating_point() or not bool(torch.isfinite(values).all()):
+        raise ValueError("Top-k values must be finite floating point tensors.")
+    # Python integer bounds avoid narrowing 2**31 to int32 (or 16384 to int16).
+    if indices.numel() and (int(indices.min()) < 0 or int(indices.max()) >= width):
+        raise ValueError(f"Top-k code index is outside 0..{width - 1}.")
+    # Sorting validates uniqueness without changing the stored (residues,k) order.
+    ordered = indices.sort(dim=1).values  # (r, k)
+    if bool((ordered[:, 1:] == ordered[:, :-1]).any()):
+        raise ValueError("Top-k indices must be unique within each residue.")
+
+
+def encode_topk_rows(
+    rows: Sequence[TopKRow], width: int, count: int, dtype: torch.dtype,
+) -> dict[str, Tensor]:
+    """Encode sparse residue rows without allocating a residue-by-codebook tensor."""
+    # rows[i].indices, .values: (r_i, k), with k = count
+    offsets = torch.zeros(len(rows) + 1, dtype=OFFSET_DTYPE)  # (n + 1,)
+    indices, values = [], []
+    for position, row in enumerate(rows):
+        if not isinstance(row, TopKRow):
+            raise TypeError("A ragged top-k feature requires TopKRow values.")
+        row_indices = row.indices.detach().to("cpu")  # (r_i, k)
+        row_values = row.values.detach().to("cpu")  # (r_i, k)
+        validate_topk(row_indices, row_values, width, count)
+        converted = row_values.to(dtype)  # (r_i, k)
+        if not bool(torch.isfinite(converted).all()):
+            raise ValueError("Top-k value conversion exceeded the stored dtype range.")
+        offsets[position + 1] = int(offsets[position]) + row_values.shape[0]
+        indices.append(row_indices.to(INDEX_DTYPE))
+        values.append(converted)
+    return {
+        "offsets": offsets,  # (n + 1,)
+        "indices": torch.cat(indices) if indices else torch.empty((0, count), dtype=INDEX_DTYPE),
+        "values": torch.cat(values) if values else torch.empty((0, count), dtype=dtype),
+    }  # offsets (n + 1,); indices and values (sum r_i, k)
+
+
+def row_count(layout: str, tensors: dict[str, Tensor]) -> int:
+    """How many sequences a segment's tensors hold."""
+
+    # tensors: (n, w) values for dense, (n + 1,) indptr or offsets otherwise
+    if layout == DENSE:
+        return int(tensors["values"].shape[0])
+    if layout == CSR:
+        return int(tensors["indptr"].shape[0]) - 1
+    if layout in (RAGGED, RAGGED_TOPK):
+        return int(tensors["offsets"].shape[0]) - 1
+    raise ValueError(f"Unknown feature layout {layout!r}; expected one of {list(LAYOUT_NAMES)}.")
+
+
+def tensor_names(layout: str, *, positions: bool) -> tuple[str, ...]:
+    """The tensors a segment of this layout holds, in a stable order."""
+
+    if layout == DENSE:
+        return ("values",)
+    if layout == CSR:
+        return ("indptr", "indices", "values", "positions") if positions else (
+            "indptr", "indices", "values",
+        )
+    if layout == RAGGED:
+        return ("offsets", "values")
+    if layout == RAGGED_TOPK:
+        return ("offsets", "indices", "values")
+    raise ValueError(f"Unknown feature layout {layout!r}; expected one of {list(LAYOUT_NAMES)}.")
+
+
+def _concatenate(parts: Sequence[Tensor], dtype: torch.dtype) -> Tensor:
+    # parts: (n_i, ...) each, with one shared trailing shape
+    if not parts:
+        return torch.empty(0, dtype=dtype)  # (0,)
+    return torch.cat(list(parts), dim=0)  # (sum n_i, ...)
+
+
+__all__ = [
+    "CSR",
+    "DENSE",
+    "DTYPE_NAMES",
+    "INDEX_DTYPE",
+    "LAYOUT_NAMES",
+    "NAME_DTYPES",
+    "OFFSET_DTYPE",
+    "POSITION_DTYPE",
+    "RAGGED",
+    "RAGGED_TOPK",
+    "VALUE_DTYPES",
+    "SparseRow",
+    "TopKRow",
+    "dtype_name",
+    "encode_csr",
+    "encode_dense",
+    "encode_ragged",
+    "encode_topk_rows",
+    "row_count",
+    "row_tensor_bytes",
+    "tensor_names",
+    "validate_topk",
+    "value_dtype",
+]

fastplms/features/reader.py

@@ -0,0 +1,546 @@
+"""Random-access reads of one feature, for loops that ask for a few rows at a time.
+
+``FeatureStore.read`` is the verified one-shot read. Each call opens a fresh index connection,
+hashes every part the request touches, and loads each of those parts whole, so a training loop that
+reads one batch per step pays for the whole store on every step. ``FeatureReader`` pays the same
+verification once per part, then serves each row by a memory-mapped slice of the part file.
+
+A reader gives the same rows as the store, in the same types, and refuses what the store refuses:
+a missing sequence raises, changed part bytes raise on first access, and independent content pins
+are honored. What it changes is the cost:
+one read-only index connection per thread, one open handle per verified part, and one slice per row.
+
+A reader belongs to a single feature directory and may be shared by threads. It pickles as the store
+it reads, so a spawned worker reopens and re-verifies only the parts it touches. Call ``verify``
+before forking to pay for verification once instead of once per worker.
+
+A reader given a ``receipt`` path remembers its verifications across processes (see ``receipts``):
+a part whose marker digests, size and modification time match the receipt is opened without hashing
+or loading it, and every part the reader does verify is recorded. A pickled reader keeps its
+receipt, so spawned workers skip the verification too. A store opened with content pins never trusts
+a receipt.
+"""
+
+from __future__ import annotations
+
+import gzip
+import json
+import os
+import sqlite3
+import sys
+import threading
+import torch
+
+from collections.abc import Iterable, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from typing import Any, cast
+from torch import Tensor
+from tqdm import tqdm
+
+from .digests import file_sha256
+from .layouts import CSR, DENSE, RAGGED_TOPK, SparseRow, TopKRow
+from .receipts import PartReceipt
+from .store import (
+    COMMIT_FILE,
+    FEATURE_FILE,
+    INDEX_FILE,
+    PART_TEMPLATE,
+    SEGMENTS_DIRECTORY,
+    FeatureStore,
+    RowAddress,
+    StoredFeature,
+    sequence_digest,
+)
+
+
+_LOOKUP_CHUNK = 512
+
+
+@dataclass(frozen=True, slots=True)
+class CsrRows:
+    """Rows of a csr feature gathered into one compressed-sparse-row block, in the order asked.
+
+    ``indptr`` is ``(n + 1,)`` int64 and names each row's span of ``indices``, ``values`` and
+    ``positions``, which are ``(nnz,)``. Wrap them in whatever sparse type the caller uses: no
+    array library is imported here.
+    """
+
+    indptr: Tensor
+    indices: Tensor
+    values: Tensor
+    positions: Tensor | None
+
+
+@dataclass(slots=True)
+class _VerifiedPart:
+    """A part that passed verification: its open handle, and its row offsets in memory.
+
+    Offsets are ``(n + 1,)`` int64 and small next to the values they index, so they stay resident.
+    Values, indices and positions are read by slice from ``handle``.
+    """
+
+    handle: Any
+    offsets: Tensor | None
+
+
+class FeatureReader:
+    """Verified random access to one feature's rows."""
+
+    def __init__(
+        self, store: FeatureStore, *, receipt: str | Path | None = None, trust_receipt: bool = True,
+    ) -> None:
+        self.store = store
+        self._receipt_path = None if receipt is None else Path(receipt)
+        self._trust_receipt = trust_receipt
+        # Pinned stores check each file against its pin, which a receipt cannot stand in for.
+        self._receipt = None
+        if self._receipt_path is not None and store._content_pins is None:
+            self._receipt = PartReceipt(
+                self._receipt_path, store.directory, store.spec.payload(), trust=trust_receipt,
+            )
+        self._lock = threading.RLock()
+        self._connections: dict[tuple[int, int], sqlite3.Connection] = {}
+        self._markers: dict[str, dict[str, Any]] = {}
+        self._parts: dict[tuple[str, int], _VerifiedPart] = {}
+        # Different parts verify concurrently; one part verifies only once.
+        self._part_locks: dict[tuple[str, int], threading.Lock] = {}
+        self._closed = False
+
+    @classmethod
+    def open(
+        cls, directory: str | Path, *, content_pins: dict[str, str] | None = None,
+    ) -> FeatureReader:
+        """A reader over an existing feature directory, which it never creates or repairs."""
+
+        return cls(FeatureStore.read_only(directory, content_pins=content_pins))
+
+    def __reduce__(self) -> tuple[Any, tuple[FeatureStore]]:
+        restore = partial(
+            FeatureReader, receipt=self._receipt_path, trust_receipt=self._trust_receipt,
+        )
+        return (restore, (self.store,))
+
+    def __enter__(self) -> FeatureReader:
+        return self
+
+    def __exit__(self, *exception: object) -> None:
+        self.close()
+
+    @property
+    def spec(self) -> StoredFeature:
+        return self.store.spec
+
+    def close(self) -> None:
+        """Save the receipt, then release every connection and part handle; reads then raise."""
+
+        if self._receipt is not None:
+            self._receipt.save()
+        with self._lock:
+            self._closed = True
+            for connection in self._connections.values():
+                connection.close()
+            self._connections.clear()
+            self._parts.clear()
+
+    # Membership --------------------------------------------------------------
+
+    def __len__(self) -> int:
+        return int(self._connection().execute("SELECT count(*) FROM rows").fetchone()[0])
+
+    def __contains__(self, sequence: str) -> bool:
+        digest = sequence_digest(sequence)
+        return digest in self._found([digest])
+
+    def missing(self, sequences: Iterable[str]) -> tuple[str, ...]:
+        """The sequences this feature lacks, in the order given, without repeats."""
+
+        wanted: dict[str, str] = {}
+        for sequence in sequences:
+            wanted.setdefault(sequence_digest(sequence), sequence)
+        present = self._found(list(wanted))
+        return tuple(sequence for digest, sequence in wanted.items() if digest not in present)
+
+    # Reading -----------------------------------------------------------------
+
+    def verify(
+        self, sequences: Sequence[str] | None = None, *, workers: int = 1,
+        progress: str | None = None,
+    ) -> int:
+        """Verify the parts holding these sequences, or every committed part, and count them.
+
+        Reading verifies lazily, so this is only for paying the cost at a moment of the caller's
+        choosing, such as before a data loader forks its workers. Verifying a part hashes its
+        bytes and loads its tensors, so ``workers`` threads verify that many parts at a time
+        (hashing and the loads release the interpreter lock). The first part to fail raises,
+        and parts not yet started are not verified. A part the receipt vouches for is opened
+        without hashing; the receipt is saved when the pass ends. ``progress`` names a progress
+        bar on stderr, one tick per part.
+        """
+
+        if workers < 1:
+            raise ValueError("verify needs at least one worker.")
+        if sequences is not None:
+            wanted = {(address.segment, address.part) for _, address in self._resolved(sequences)}
+        else:
+            wanted = {
+                (fingerprint, int(part["part"]))
+                for fingerprint in self._segment_names()
+                for part in self._marker(fingerprint)["parts"]
+            }
+        ordered = sorted(wanted)
+        bar = tqdm(
+            total=len(ordered), desc=progress, unit="part", file=sys.stderr, mininterval=10.0,
+            disable=progress is None,
+        )
+
+        def verified(segment: str, number: int) -> None:
+            self._part(segment, number)
+            bar.update()
+
+        try:
+            if workers == 1 or len(ordered) < 2:
+                for segment, number in ordered:
+                    verified(segment, number)
+                return len(wanted)
+            with ThreadPoolExecutor(max_workers=workers) as pool:
+                futures = [pool.submit(verified, segment, number) for segment, number in ordered]
+                try:
+                    for future in futures:
+                        future.result()
+                except BaseException:
+                    for future in futures:
+                        future.cancel()
+                    raise
+            return len(wanted)
+        finally:
+            bar.close()
+            if self._receipt is not None:
+                self._receipt.save()
+
+    def addresses(self, sequences: Sequence[str]) -> list[RowAddress]:
+        """Each sequence's address, in the order asked, from the index and its commit markers alone.
+
+        This reads no part, so it costs index lookups, not bytes. It raises ``KeyError`` for a
+        sequence the feature lacks and ``ValueError`` when the index and the commit marker disagree
+        about a row. Nothing here verifies the part holding the row: ``verify`` or a read does.
+        """
+
+        return [address for _, address in self._resolved(sequences)]
+
+    def content_pins(self, sequences: Sequence[str], *, workers: int = 1) -> dict[str, str]:
+        """Pin freshly hashed selection bytes after reusing this reader's layout verification.
+
+        Every selected file is hashed again, including when a verification receipt was trusted.
+        Cached part handles avoid a second whole-part tensor load. Marker changes and even
+        data rewrites preserving size and modification time fail before these pins are returned.
+        """
+        self.verify(sequences, workers=workers)
+        feature = self.store.directory / FEATURE_FILE
+        if StoredFeature.from_payload(json.loads(feature.read_text(encoding="utf-8"))) != self.spec:
+            raise ValueError("Feature descriptor changed while pinning a selection.")
+        pins = {FEATURE_FILE: file_sha256(feature)}
+        selected = {(address.segment, address.part) for _, address in self._resolved(sequences)}
+        markers: set[str] = set()
+        for segment, number in sorted(selected):
+            payload = self._marker(segment)
+            prefix = f"{SEGMENTS_DIRECTORY}/{segment}"
+            if segment not in markers:
+                marker = self.store.directory / prefix / COMMIT_FILE
+                if json.loads(marker.read_text(encoding="utf-8")) != payload:
+                    raise ValueError("Feature commit marker changed while pinning a selection.")
+                pins[f"{prefix}/{COMMIT_FILE}"] = file_sha256(marker)
+                markers.add(segment)
+            part = payload["parts"][number]
+            pins[f"{prefix}/{PART_TEMPLATE.format(number)}"] = part["sha256"]
+            sidecar = part.get("row_metadata")
+            if sidecar is not None:
+                pins[f"{prefix}/{sidecar['file']}"] = sidecar["sha256"]
+
+        def check(relative: str) -> None:
+            actual = file_sha256(self.store.directory / relative)
+            if actual != pins[relative]:
+                raise ValueError("Feature bytes changed while pinning a selection.")
+            self.store._check_pin(relative, actual)
+
+        if workers == 1:
+            for relative in pins:
+                check(relative)
+        else:
+            with ThreadPoolExecutor(max_workers=workers) as pool:
+                list(pool.map(check, pins))
+        return pins
+
+    def residue_counts(self, sequences: Sequence[str]) -> list[int]:
+        """Each sequence's stored residue count, which a length-bucketed reader batches by."""
+
+        return [address.residues for _, address in self._resolved(sequences, verify=True)]
+
+    def row_metadata(self, sequences: Sequence[str]) -> list[dict[str, Any]]:
+        """Return row identities using the same immutable-part verification as feature reads."""
+        identities = []
+        parts: dict[tuple[str, int], list[dict[str, Any]]] = {}
+        for _, address in self._resolved(sequences):
+            self._part(address.segment, address.part)
+            key = (address.segment, address.part)
+            if key not in parts:
+                committed = self._marker(address.segment)["parts"][address.part]
+                sidecar = committed.get("row_metadata")
+                if not isinstance(sidecar, dict):
+                    raise ValueError(
+                        "Stored row identities are unavailable; recompute this legacy feature."
+                    )
+                path = self.store.directory / SEGMENTS_DIRECTORY / address.segment / sidecar["file"]
+                records = json.loads(gzip.decompress(path.read_bytes()))
+                if (not isinstance(records, list) or len(records) != len(committed["digests"])
+                        or not all(isinstance(record, dict) for record in records)):
+                    raise ValueError("Stored row metadata does not match the committed part rows.")
+                parts[key] = records
+            identities.append(dict(parts[key][address.row]))
+        return identities
+
+    def read(self, sequences: Sequence[str]) -> list[Tensor]:
+        """Each sequence's feature as a tensor, in the order asked, as ``FeatureStore.read`` gives.
+
+        A dense feature gives ``(w,)``, a ragged one ``(r_i, d)``, and a csr one a densified
+        ``(w,)`` row. ``ragged_topk`` requires ``read_topk``.
+        """
+
+        spec = self.spec
+        if spec.layout == CSR:
+            # Returns b vectors, each (w,).
+            return [row.to_dense(spec.width) for row in self.read_sparse(sequences)]  # (w,) per row
+        if spec.layout == RAGGED_TOPK:
+            raise ValueError(
+                "Use read_topk for sparse residue codes; implicit densification is refused."
+            )
+        rows: list[Tensor] = []
+        parts: dict[tuple[str, int], tuple[_VerifiedPart, Tensor]] = {}
+        for _, address in self._resolved(sequences):
+            key = (address.segment, address.part)
+            if key not in parts:
+                part = self._part(*key)
+                # One mapped tensor per touched part, not one safetensors wrapper per row.
+                parts[key] = (part, part.handle.get_tensor("values"))
+            part, values = parts[key]  # (n_part, w) dense or (r_part, d) ragged
+            if spec.layout == DENSE:
+                rows.append(values[address.row].clone())  # (w,)
+            else:
+                start, stop = self._span(part, address)
+                rows.append(values[start:stop].clone())  # (r_i, d)
+        return rows  # b tensors, each (w,) for dense or (r_i, d) for ragged
+
+    def read_sparse(self, sequences: Sequence[str]) -> list[SparseRow]:
+        """Each sequence's compressed row, in the order asked. Only for a csr feature."""
+
+        if self.spec.layout != CSR:
+            raise ValueError(f"read_sparse needs a csr feature; this one is {self.spec.layout}.")
+        rows: list[SparseRow] = []
+        parts: dict[tuple[str, int], tuple[_VerifiedPart, Tensor, Tensor, Tensor | None]] = {}
+        for _, address in self._resolved(sequences):
+            key = (address.segment, address.part)
+            if key not in parts:
+                part = self._part(*key)
+                parts[key] = (
+                    part, part.handle.get_tensor("indices"), part.handle.get_tensor("values"),
+                    part.handle.get_tensor("positions") if self.spec.positions else None,
+                )
+            part, indices, values, positions = parts[key]  # (nnz_part,) per tensor
+            start, stop = self._span(part, address)
+            rows.append(SparseRow(
+                indices=indices[start:stop].clone(),  # (nnz_i,)
+                values=values[start:stop].clone(),  # (nnz_i,)
+                positions=(
+                    positions[start:stop].clone() if positions is not None else None  # (nnz_i,)
+                ),
+            ))
+        return rows
+
+    def read_csr(self, sequences: Sequence[str]) -> CsrRows:
+        """These sequences' rows as one compressed-sparse-row block, in the order asked.
+
+        A sequence asked for twice appears twice. This is the read for a caller that wants a matrix,
+        such as a design matrix for a gradient-boosted model, instead of one row at a time.
+        """
+
+        spec = self.spec
+        rows = self.read_sparse(sequences)
+        counts = torch.tensor([row.indices.numel() for row in rows], dtype=torch.int64)  # (n,)
+        indptr = torch.zeros(len(rows) + 1, dtype=torch.int64)  # (n + 1,)
+        torch.cumsum(counts, dim=0, out=indptr[1:])
+
+        def joined(tensors: list[Tensor], dtype: torch.dtype) -> Tensor:
+            # tensors: (nnz_i,) per input vector.
+            return torch.cat(tensors) if tensors else torch.empty(0, dtype=dtype)  # (nnz,)
+
+        return CsrRows(
+            indptr=indptr,
+            indices=joined([row.indices for row in rows], torch.int32),  # (nnz,)
+            values=joined([row.values for row in rows], spec.dtype),  # (nnz,)
+            positions=(
+                joined([cast(Tensor, row.positions) for row in rows], torch.int16)  # (nnz,)
+                if spec.positions else None
+            ),
+        )
+
+    def read_topk(self, sequences: Sequence[str]) -> list[TopKRow]:
+        """Each sequence's sparse ``(r_i, k)`` residue codes, in the order asked."""
+
+        if self.spec.layout != RAGGED_TOPK:
+            raise ValueError(
+                f"read_topk needs a ragged_topk feature; this one is {self.spec.layout}."
+            )
+        rows: list[TopKRow] = []
+        for _, address in self._resolved(sequences):
+            part = self._part(address.segment, address.part)
+            start, stop = self._span(part, address)
+            rows.append(TopKRow(
+                cast(Tensor, part.handle.get_slice("indices")[start:stop]),  # (r_i, k)
+                cast(Tensor, part.handle.get_slice("values")[start:stop]),  # (r_i, k)
+            ))
+        return rows
+
+    # Internals ---------------------------------------------------------------
+
+    def _connection(self) -> sqlite3.Connection:
+        """This thread's read-only index connection, opened on first use.
+
+        Connections are keyed by process and thread, so a forked worker or a prefetch thread opens
+        its own and never touches one that another owner holds.
+        """
+
+        owner = (os.getpid(), threading.get_ident())
+        with self._lock:
+            if self._closed:
+                raise ValueError("This feature reader is closed.")
+            connection = self._connections.get(owner)
+            if connection is None:
+                database = (self.store.directory / INDEX_FILE).resolve()
+                connection = sqlite3.connect(
+                    database.as_uri() + "?mode=ro", uri=True, timeout=30,
+                    check_same_thread=False,
+                )
+                self._connections[owner] = connection
+            return connection
+
+    def _found(self, digests: Sequence[str]) -> dict[str, RowAddress]:
+        found: dict[str, RowAddress] = {}
+        connection = self._connection()
+        for start in range(0, len(digests), _LOOKUP_CHUNK):
+            chunk = list(digests[start:start + _LOOKUP_CHUNK])
+            marks = ",".join("?" * len(chunk))
+            # fetchall ends the read transaction so writers can commit.
+            for digest, segment, part, row, residues in connection.execute(
+                f"SELECT digest, segment, part, row, residues FROM rows WHERE digest IN ({marks})",
+                chunk,
+            ).fetchall():
+                found[digest] = RowAddress(segment, part, row, residues)
+        return found
+
+    def _resolved(
+        self, sequences: Sequence[str], *, verify: bool = False,
+    ) -> list[tuple[str, RowAddress]]:
+        """Each sequence with its address, checked against the commit marker that owns it."""
+
+        self.store._check_pin(FEATURE_FILE)
+        digests = [sequence_digest(sequence) for sequence in sequences]
+        found = self._found(list(dict.fromkeys(digests)))
+        resolved: list[tuple[str, RowAddress]] = []
+        for sequence, digest in zip(sequences, digests, strict=True):
+            address = found.get(digest)
+            if address is None:
+                raise KeyError(
+                    f"Feature {self.spec.key!r} has no row for a sequence of "
+                    f"{len(sequence)} residues (sha256 {digest[:12]}). "
+                    "Embed it first, or call missing() before reading."
+                )
+            payload = self._marker(address.segment)
+            if not 0 <= address.part < len(payload["parts"]):
+                raise ValueError("Index part does not name one committed feature part.")
+            part = payload["parts"][address.part]
+            if (not 0 <= address.row < len(part["digests"])
+                    or part["digests"][address.row] != digest
+                    or part["residues"][address.row] != address.residues):
+                raise ValueError("Feature index and committed sequence row disagree.")
+            if verify:
+                self._part(address.segment, address.part)
+            resolved.append((sequence, address))
+        return resolved
+
+    def _marker(self, segment: str) -> dict[str, Any]:
+        """A validated commit marker, read once per reader."""
+
+        with self._lock:
+            payload = self._markers.get(segment)
+            if payload is None:
+                payload = self.store._segment_payload(segment)
+                self._markers[segment] = payload
+            return payload
+
+    def _segment_names(self) -> list[str]:
+        directory = self.store.directory / SEGMENTS_DIRECTORY
+        return sorted(marker.parent.name for marker in directory.glob(f"*/{COMMIT_FILE}"))
+
+    def _part(self, segment: str, number: int) -> _VerifiedPart:
+        """A part's open handle, verified the first time and kept.
+
+        Verification is the store's own: checksum against the marker, physical layout, and the row
+        identity sidecar. The tensors it loads are dropped afterwards except the offsets.
+        """
+
+        key = (segment, number)
+        with self._lock:
+            if self._closed:
+                raise ValueError("This feature reader is closed.")
+            opened = self._parts.get(key)
+            if opened is not None:
+                return opened
+            part_lock = self._part_locks.setdefault(key, threading.Lock())
+        with part_lock:
+            # Verify outside the reader lock so other parts can progress.
+            with self._lock:
+                opened = self._parts.get(key)
+            if opened is not None:
+                return opened
+            from safetensors import safe_open
+
+            committed = self._marker(segment)["parts"][number]
+            path = str(self.store._part_path(segment, number))
+            offsets_name = "indptr" if self.spec.layout == CSR else "offsets"
+            receipt = self._receipt
+            described = None if receipt is None else receipt.describe(segment, number, committed)
+            if (receipt is not None and described is not None
+                    and receipt.trusts(segment, number, described)):
+                # An earlier full verification passed these digests on files of this size and time.
+                handle = safe_open(path, framework="pt", device="cpu")
+                names = set(handle.keys())
+                offsets = handle.get_tensor(offsets_name) if offsets_name in names else None
+            else:
+                tensors = self.store._verified_part(segment, committed)
+                offsets = tensors.get(offsets_name)
+                del tensors
+                handle = safe_open(path, framework="pt", device="cpu")
+                if receipt is not None and described is not None:
+                    if receipt.describe(segment, number, committed) != described:
+                        raise ValueError(
+                            f"A feature part changed during verification: {segment}/{number}."
+                        )
+                    receipt.record(segment, number, described)
+            opened = _VerifiedPart(handle=handle, offsets=offsets)
+            with self._lock:
+                if self._closed:
+                    raise ValueError("This feature reader is closed.")
+                self._parts[key] = opened
+            return opened
+
+    @staticmethod
+    def _span(part: _VerifiedPart, address: RowAddress) -> tuple[int, int]:
+        """A row's ``[start, stop)`` span in its part's flattened values."""
+
+        assert part.offsets is not None, "Only sparse and ragged layouts carry row offsets."
+        return int(part.offsets[address.row]), int(part.offsets[address.row + 1])
+
+
+__all__ = ["FeatureReader"]

fastplms/features/receipts.py

@@ -0,0 +1,125 @@
+"""A record of which feature parts a full verification has passed, so a later reader can skip it.
+
+Verifying a part hashes its bytes and loads its tensors, which costs one pass over the file. A store
+of hundreds of gigabytes read by many processes would otherwise pay that pass in every process, on
+every open. A receipt keeps the result for one feature directory: for each part, the digests the
+commit marker states, and the size and modification time of the part file and of its row-identity
+sidecar when they were verified. A reader trusts a part only when the marker states the same digests
+and both files still have the recorded size and modification time. Any other part is verified in
+full.
+
+A receipt is a cache of one machine's verification, never a proof: a file rewritten to the same size
+and the same modification time passes it. ``FeatureReader.verify`` on a reader built with
+``trust_receipt=False`` hashes every byte and refreshes the receipt, and is the check for that case.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import threading
+import warnings
+
+from collections.abc import Mapping
+from contextlib import suppress
+from pathlib import Path
+from typing import Any
+
+from .digests import json_sha256
+from .store import PART_TEMPLATE, SEGMENTS_DIRECTORY
+
+
+SCHEMA = "feature_part_receipt_v1"
+
+
+class PartReceipt:
+    """The verified parts of one feature directory, read from and saved to one JSON file."""
+
+    def __init__(
+        self, path: str | Path, directory: Path, spec_payload: Mapping[str, Any],
+        *, trust: bool = True,
+    ) -> None:
+        self.path = Path(path)
+        self._directory = directory
+        self._feature = spec_payload.get("key")
+        self._descriptor_sha256 = json_sha256(spec_payload)
+        self._trusted = self._load() if trust else {}
+        self._pending: dict[str, dict[str, Any]] = {}
+        self._lock = threading.Lock()
+
+    def _load(self) -> dict[str, dict[str, Any]]:
+        """The recorded parts, or none without a receipt or when it names another descriptor."""
+        try:
+            document = json.loads(self.path.read_text(encoding="utf-8"))
+        except FileNotFoundError:
+            return {}
+        except (OSError, ValueError) as error:
+            warnings.warn(
+                f"Ignoring the unreadable verification receipt {self.path}: {error}",
+                RuntimeWarning, stacklevel=3,
+            )
+            return {}
+        if (not isinstance(document, dict) or document.get("schema") != SCHEMA
+                or document.get("descriptor_sha256") != self._descriptor_sha256
+                or not isinstance(document.get("parts"), dict)):
+            return {}
+        return document["parts"]
+
+    def describe(self, segment: str, number: int, committed: Mapping[str, Any]) -> dict[str, Any]:
+        """What verifying this part vouches for: the marker's digests and the files' stat."""
+        base = self._directory / SEGMENTS_DIRECTORY / segment
+        part = (base / PART_TEMPLATE.format(number)).stat()
+        sidecar = committed.get("row_metadata")
+        rows = None
+        if isinstance(sidecar, dict):
+            stat = (base / str(sidecar["file"])).stat()
+            rows = {
+                "file": sidecar["file"], "sha256": sidecar["sha256"],
+                "size": stat.st_size, "mtime_ns": stat.st_mtime_ns,
+            }
+        return {
+            "sha256": committed["sha256"], "size": part.st_size, "mtime_ns": part.st_mtime_ns,
+            "rows": rows,
+        }
+
+    def trusts(self, segment: str, number: int, described: Mapping[str, Any]) -> bool:
+        """Whether an earlier verification passed these digests, sizes and modification times."""
+        return self._trusted.get(f"{segment}/{number}") == described
+
+    def record(self, segment: str, number: int, described: Mapping[str, Any]) -> None:
+        with self._lock:
+            self._pending[f"{segment}/{number}"] = dict(described)
+
+    def save(self) -> None:
+        """Merge the parts verified since the last save into the file, atomically.
+
+        Concurrent savers can drop each other's entries, which only costs a later verification. A
+        location that cannot be written warns, and this process keeps what it verified in memory.
+        """
+        with self._lock:
+            if not self._pending:
+                return
+            document = {
+                "schema": SCHEMA, "feature": self._feature,
+                "descriptor_sha256": self._descriptor_sha256,
+                "parts": {**self._load(), **self._pending},
+            }
+            owner = f"{os.getpid()}.{threading.get_ident()}"
+            temporary = self.path.with_name(f"{self.path.name}.{owner}.writing")
+            try:
+                self.path.parent.mkdir(parents=True, exist_ok=True)
+                temporary.write_text(json.dumps(document, sort_keys=True), encoding="utf-8")
+                temporary.replace(self.path)
+            except OSError as error:
+                with suppress(OSError):
+                    temporary.unlink(missing_ok=True)
+                warnings.warn(
+                    f"Could not save the verification receipt {self.path}: {error}. "
+                    "Every open of this feature verifies its parts again.",
+                    RuntimeWarning, stacklevel=2,
+                )
+            self._trusted = {**self._trusted, **self._pending}
+            self._pending.clear()
+
+
+__all__ = ["SCHEMA", "PartReceipt"]

fastplms/features/store.py

@@ -0,0 +1,1567 @@
+"""The FastPLMs feature store: one directory per feature key, read by sequence.
+
+A store answers one question: what is this feature of this sequence, and does it already exist?
+That makes a run embed only what is missing, and lets one embedding pass serve every head trained
+afterwards.
+
+Layout on disk, under one root that holds many features::
+
+    <root>/<key>/
+        feature.json                       the key's descriptor, layout, width, and dtype
+        index.sqlite                       sequence sha256 -> segment, part, row, residue count
+        segments/<fingerprint>/
+            part-00000.safetensors         immutable, memory-mappable, one per append
+            part-00001.safetensors
+            run.json                       the commit marker, written last
+
+**A segment is immutable and committed once.** Parts appear as a run streams, and nothing reads
+them until ``run.json`` names them; a run that dies leaves a directory the index ignores and
+``sweep`` removes. Nothing is ever rewritten, so a reader never sees a half-written feature and two
+runs never race over one file.
+
+**The index is a cache of the commit markers, not the record.** Every indexed row can be rebuilt
+from the committed segments, which ``reindex`` does, so a lost or corrupt ``index.sqlite`` costs a
+scan rather than the features.
+
+**A sequence is identified by the SHA-256 of its exact UTF-8 bytes**, so two callers agree without
+coordinating, and a sequence that differs by one residue is a different row. The store keeps the
+digest and the residue count, never the sequence text: a caller that has the sequences can always
+recompute the digest, and storing millions of them again would cost more than the features.
+
+The key belongs to ``foundry.embedding.feature_key``, which composes the model, its revision, the
+sparse autoencoder, the layer, the pooling, the dtype, and the residue limit into one filename-safe
+name. This module never invents a key; it stores the name and the descriptor it is given and
+refuses a second, different descriptor under the same name. FastPLMs ships without foundry, so the
+store takes the name as a string rather than importing the key.
+"""
+
+from __future__ import annotations
+
+import gzip
+import hashlib
+import json
+import os
+import re
+import sqlite3
+import threading
+import torch
+
+from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
+from contextlib import AbstractContextManager, ExitStack, contextmanager
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from itertools import pairwise
+from pathlib import Path
+from typing import Any, cast
+from torch import Tensor
+
+from .digests import file_sha256, json_sha256
+from .json_files import indented_json
+from .layouts import (
+    CSR,
+    DENSE,
+    LAYOUT_NAMES,
+    RAGGED,
+    RAGGED_TOPK,
+    SparseRow,
+    TopKRow,
+    dtype_name,
+    encode_csr,
+    encode_dense,
+    encode_ragged,
+    encode_topk_rows,
+    row_count,
+    row_tensor_bytes,
+    tensor_names,
+    validate_topk,
+    value_dtype,
+)
+from .transactions import file_lock, flush_and_evict, publish_file, sync_directory
+
+
+FORMAT = "fastplms-feature-store-v1"
+FEATURE_FILE = "feature.json"
+INDEX_FILE = "index.sqlite"
+SEGMENTS_DIRECTORY = "segments"
+COMMIT_FILE = "run.json"
+PART_TEMPLATE = "part-{:05d}.safetensors"
+
+# Descriptor schemas that carry a complete scientific contract: v1 keeps residues only, v2 keeps CLS and EOS too,
+# and v3 keeps v2's rows computed under a pinned embedding profile.
+COMPLETE_SCHEMAS = frozenset({"feature_spec_v1", "feature_spec_v2", "feature_spec_v3"})
+
+_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,255}")
+_PART = re.compile(r"part-(\d{5})\.safetensors")
+_SHA256 = re.compile(r"[0-9a-f]{64}")
+_WRITER_LEASE = object()
+
+
+def sequence_digest(sequence: str) -> str:
+    """The SHA-256 of a sequence's exact UTF-8 bytes, which is its row identity."""
+
+    if not isinstance(sequence, str) or not sequence:
+        raise ValueError("A sequence must be a non-empty string.")
+    return hashlib.sha256(sequence.encode("utf-8")).hexdigest()
+
+
+def partition_sequences(sequences: Iterable[str], *, shard: int, shards: int) -> tuple[str, ...]:
+    """Assign exact unique sequences by SHA-256 modulo shard count, preserving input order.
+
+    Every worker must receive the same input inventory and shard count. Assignment does not
+    depend on Python's randomized hash or process/device identity.
+    """
+    if type(shards) is not int or shards < 1 or type(shard) is not int or not 0 <= shard < shards:
+        raise ValueError("Require a positive shard count and 0 <= shard < shards.")
+    unique = dict.fromkeys(sequences)
+    return tuple(
+        sequence for sequence in unique if int(sequence_digest(sequence), 16) % shards == shard
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class StoredFeature:
+    """What one feature is: its key, how its rows are laid out, and what they hold.
+
+    ``width`` is the vector width for ``dense``, the codebook size for ``csr``, and the hidden
+    width for ``ragged``. For ``ragged_topk``, ``width`` is the codebook size and ``sparse_count``
+    is the number of retained codes per residue. ``positions`` says whether ``csr`` rows carry
+    each entry's argmax residue. ``descriptor`` is the plain data the key was computed from,
+    so a store can say what it holds without the caller that made it.
+    """
+
+    key: str
+    layout: str
+    width: int
+    dtype: torch.dtype
+    positions: bool = False
+    descriptor: Mapping[str, Any] = field(default_factory=dict)
+    sparse_count: int | None = None
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.key, str) or not _NAME.fullmatch(self.key):
+            raise ValueError(
+                "A feature key must be a filename-safe name, as "
+                "foundry.embedding.feature_key(...).name returns; received "
+                f"{self.key!r}."
+            )
+        if self.layout not in LAYOUT_NAMES:
+            raise ValueError(
+                f"layout must be one of {list(LAYOUT_NAMES)}; received {self.layout!r}."
+            )
+        if not isinstance(self.width, int) or isinstance(self.width, bool) or self.width <= 0:
+            raise ValueError(f"width must be a positive integer; received {self.width!r}.")
+        dtype_name(self.dtype)
+        if self.positions and self.layout != CSR:
+            raise ValueError("Only a csr feature stores argmax positions.")
+        if self.layout == RAGGED_TOPK:
+            if type(self.sparse_count) is not int or not 1 <= self.sparse_count <= self.width:
+                raise ValueError(
+                    "A ragged top-k feature requires integer sparse_count in 1..width."
+                )
+            if self.width > torch.iinfo(torch.int32).max + 1:
+                raise ValueError("Top-k codebook exceeds the stored int32 index range.")
+        elif self.sparse_count is not None:
+            raise ValueError("Only a ragged top-k feature declares sparse_count.")
+        if not isinstance(self.descriptor, Mapping):
+            raise TypeError("descriptor must be a mapping of plain data.")
+        object.__setattr__(self, "descriptor", json.loads(json.dumps(dict(self.descriptor))))
+
+    def payload(self) -> dict[str, Any]:
+        """``feature.json``'s content."""
+
+        payload = {
+            "format": FORMAT,
+            "key": self.key,
+            "layout": self.layout,
+            "width": self.width,
+            "dtype": dtype_name(self.dtype),
+            "positions": self.positions,
+            "descriptor": dict(self.descriptor),
+        }
+        if self.sparse_count is not None:
+            payload["sparse_count"] = self.sparse_count
+        return payload
+
+    @classmethod
+    def from_payload(cls, payload: Mapping[str, Any]) -> StoredFeature:
+        """The spec a ``feature.json`` describes."""
+
+        if payload.get("format") != FORMAT:
+            raise ValueError(
+                f"Not a {FORMAT} feature directory; its format is {payload.get('format')!r}."
+            )
+        return cls(
+            key=str(payload["key"]),
+            layout=str(payload["layout"]),
+            width=int(payload["width"]),
+            dtype=value_dtype(str(payload["dtype"])),
+            positions=bool(payload.get("positions", False)),
+            descriptor=payload.get("descriptor") or {},
+            sparse_count=payload.get("sparse_count"),
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class RowAddress:
+    """Where one sequence's row sits: which segment, which part, which row of it."""
+
+    segment: str
+    part: int
+    row: int
+    residues: int
+
+
+@dataclass(frozen=True, slots=True)
+class SegmentReceipt:
+    """What one committed segment holds, as ``run.json`` records it."""
+
+    fingerprint: str
+    parts: tuple[int, ...]
+    rows: int
+    committed_at: str
+    metadata: Mapping[str, Any]
+
+
+class FeatureStore:
+    """One feature's directory, opened for reading and appending."""
+
+    def __init__(
+        self, directory: Path, spec: StoredFeature, *, read_only: bool = False,
+        content_pins: Mapping[str, str] | None = None, deep_verify: bool = True,
+        partial: bool = False,
+    ) -> None:
+        self.directory = directory.resolve()
+        self.spec = spec
+        self._read_only = read_only
+        # Deep verification re-reads every committed part whenever the index is recovered, which opening
+        # and every commit do. A store of terabytes opens with deep_verify=False: only segments the index
+        # lacks are read and indexed, and each part is still verified when a reader first touches it.
+        self._deep_verify = deep_verify
+        # A partial copy holds some of each segment's committed parts, as a fetch of a few rows from a
+        # published store leaves it. Its index lists the rows of the parts present, so a row of an absent
+        # part reads as missing, a part fetched later is indexed on the next open, and it takes no new
+        # segment. Absent parts cannot be verified, so it opens without deep verification.
+        if partial and deep_verify:
+            raise ValueError("A partial copy opens with deep_verify=False; its absent parts cannot be read.")
+        self._partial = partial
+        self._content_pins = None if content_pins is None else dict(content_pins)
+        if self._content_pins is not None:
+            if not read_only or FEATURE_FILE not in self._content_pins:
+                raise ValueError("Pinned feature handles must be read-only and pin feature.json.")
+            for relative, digest in self._content_pins.items():
+                components = relative.split("/") if isinstance(relative, str) else []
+                valid_path = relative == FEATURE_FILE or (
+                    len(components) == 3 and components[0] == SEGMENTS_DIRECTORY
+                    and _NAME.fullmatch(components[1])
+                    and (components[2] == COMMIT_FILE or re.fullmatch(
+                        r"part-\d{5,}\.(safetensors|rows\.json\.gz)", components[2],
+                    ))
+                )
+                if not valid_path or not isinstance(digest, str) or not _SHA256.fullmatch(digest):
+                    raise ValueError(
+                        "Feature content pins must name immutable store files and SHA-256 digests."
+                    )
+
+    # Opening -----------------------------------------------------------------
+
+    @classmethod
+    def open(
+        cls, root: str | Path, spec: StoredFeature, *, deep_verify: bool = True, partial: bool = False,
+    ) -> FeatureStore:
+        """Open the store for ``spec`` under ``root``, creating it when it does not exist.
+
+        An existing directory must describe exactly this spec. A mismatch is a caller using one key
+        for two different features, which would serve one project another project's numbers, so it
+        raises rather than migrating. ``deep_verify=False`` is for stores too large to re-read on
+        every open and commit, and ``partial=True`` for a copy holding only some committed parts,
+        which must already exist: see the constructor.
+        """
+
+        directory = Path(root) / spec.key
+        recorded = directory / FEATURE_FILE
+        if partial and not recorded.exists():
+            raise FileNotFoundError(f"{directory} holds no {FEATURE_FILE}, so it is no copy of a store.")
+        store = cls(directory, spec, deep_verify=deep_verify, partial=partial)
+        with store._write_lock():
+            if recorded.exists():
+                existing = StoredFeature.from_payload(
+                    json.loads(recorded.read_text(encoding="utf-8")))
+                if existing != spec:
+                    raise ValueError(
+                        f"{directory} already holds a different feature under key {spec.key!r}.\n"
+                        f"  stored:    {existing.payload()}\n"
+                        f"  requested: {spec.payload()}"
+                    )
+            else:
+                (directory / SEGMENTS_DIRECTORY).mkdir(parents=True, exist_ok=True)
+                _write_json_atomically(recorded, spec.payload())
+            store._ensure_index_locked()
+        return store
+
+    @classmethod
+    def read_only(
+        cls, directory: str | Path, *, content_pins: Mapping[str, str] | None = None,
+    ) -> FeatureStore:
+        """Open an existing descriptor without creating or repairing any file.
+
+        Each query opens its own SQLite connection in read-only mode and closes it before
+        returning. A missing or corrupt index fails when queried, without being repaired.
+        """
+
+        path = Path(directory)
+        payload = json.loads((path / FEATURE_FILE).read_text(encoding="utf-8"))
+        store = cls(
+            path, StoredFeature.from_payload(payload), read_only=True, content_pins=content_pins,
+        )
+        store._check_pin(FEATURE_FILE)
+        return store
+
+    # Membership --------------------------------------------------------------
+
+    def __len__(self) -> int:
+        with self._connect() as connection:
+            return int(connection.execute("SELECT count(*) FROM rows").fetchone()[0])
+
+    def address(self, sequence: str) -> RowAddress | None:
+        """Where this sequence's row is, or None when the store lacks it."""
+
+        with self._connect() as connection:
+            found = connection.execute(
+                "SELECT segment, part, row, residues FROM rows WHERE digest = ?",
+                (sequence_digest(sequence),),
+            ).fetchone()
+        return None if found is None else RowAddress(*found)
+
+    def missing(self, sequences: Iterable[str]) -> tuple[str, ...]:
+        """The sequences this store lacks, in the order given, without repeats.
+
+        This is the call that makes a run embed only what is new.
+        """
+
+        wanted: dict[str, str] = {}
+        for sequence in sequences:
+            wanted.setdefault(sequence_digest(sequence), sequence)
+        if not wanted:
+            return ()
+        present = self.present_digests(list(wanted))
+        return tuple(sequence for digest, sequence in wanted.items() if digest not in present)
+
+    def present_digests(self, digests: Sequence[str]) -> set[str]:
+        """The digests among ``digests`` that this store already holds, from one index connection."""
+
+        present: set[str] = set()
+        with self._connect() as connection:
+            for chunk in _chunks(list(digests), 512):
+                marks = ",".join("?" * len(chunk))
+                present.update(
+                    row[0]
+                    for row in connection.execute(
+                        f"SELECT digest FROM rows WHERE digest IN ({marks})", chunk
+                    )
+                )
+        return present
+
+    def segments(self) -> tuple[SegmentReceipt, ...]:
+        """Every committed segment, oldest commit first, then by fingerprint.
+
+        Commit times have one-second resolution, so the fingerprint breaks the tie and the order is
+        total.
+        """
+
+        receipts = [_receipt(payload) for payload in self._verified_segments()]
+        return tuple(sorted(receipts, key=lambda receipt: (receipt.committed_at, receipt.fingerprint)))
+
+    # Writing -----------------------------------------------------------------
+
+    @contextmanager
+    def segment(
+        self, fingerprint: str, metadata: Mapping[str, Any] | None = None,
+        *, before_commit: Callable[[], None] | None = None, verify_staged: bool = True,
+    ) -> Iterator[SegmentWriter]:
+        """Open a new segment for one embedding run, committed on a clean exit.
+
+        ``verify_staged`` makes the commit re-read and re-hash every staged file. A caller that wrote
+        large parts through ``append_packed`` and hashed them as it wrote may pass False: the commit
+        then compares each file's size and modification time with what the write recorded, which
+        catches an appended or rewritten file without reading the payload a second time.
+
+        ``fingerprint`` identifies the run that produced these rows, and is FastPLMs' own run
+        fingerprint in a real run. A committed segment of that name already existing is an error:
+        the same run producing the same rows twice means one of the two is not what it claims.
+
+        A body that writes nothing leaves nothing behind, and a body that raises leaves the segment
+        uncommitted for ``sweep``.
+        """
+
+        self._require_writable()
+        if self._partial:
+            raise PermissionError("A partial copy takes no new segment; embed into a store of its own.")
+        if not isinstance(fingerprint, str) or not _NAME.fullmatch(fingerprint):
+            raise ValueError(f"A segment fingerprint must be a filename-safe name; received {fingerprint!r}.")
+        if self.spec.descriptor.get("schema") in COMPLETE_SCHEMAS and before_commit is None:
+            raise ValueError(
+                "A complete feature contract requires a pre-commit validation callback."
+            )
+        path = self.directory / SEGMENTS_DIRECTORY / fingerprint
+        with file_lock(self._segment_lock_path(fingerprint), wait=False):
+            with self._write_lock():
+                if (path / COMMIT_FILE).exists():
+                    raise FileExistsError(
+                        f"Segment {fingerprint!r} is already committed in {self.directory}."
+                    )
+                # Only the owner of this segment lock can discard a crashed attempt.
+                if path.exists():
+                    _discard_segment(path)
+                path.mkdir(parents=True)
+                sync_directory(path.parent)
+            writer = SegmentWriter(
+                self, path, fingerprint, dict(metadata or {}), before_commit, lease=_WRITER_LEASE,
+                verify_staged=verify_staged,
+            )
+            try:
+                yield writer
+                if not writer.committed and not writer.closed:
+                    if writer.parts:
+                        writer.commit()
+                    else:
+                        writer.abandon()
+            finally:
+                writer.closed = True
+
+    def sweep(self) -> tuple[str, ...]:
+        """Delete every uncommitted segment directory, and name what was deleted.
+
+        An uncommitted segment is the remains of a run that died. Nothing reads it.
+        """
+
+        self._require_writable()
+        removed: list[str] = []
+        for path in sorted((self.directory / SEGMENTS_DIRECTORY).iterdir()):
+            if not path.is_dir() or not _NAME.fullmatch(path.name):
+                continue
+            with ExitStack() as stack:
+                try:
+                    stack.enter_context(file_lock(self._segment_lock_path(path.name), wait=False))
+                except BlockingIOError:
+                    continue
+                with self._write_lock():
+                    if path.exists() and not (path / COMMIT_FILE).exists():
+                        _discard_segment(path)
+                        removed.append(path.name)
+        return tuple(removed)
+
+    def reindex(self) -> int:
+        """Rebuild the index from the committed segments, and return the row count.
+
+        The index is a cache, so this is the repair when it is lost or doubted. It reads each
+        part's row count from its tensors and each part's digests from the commit marker.
+        """
+
+        self._require_writable()
+        with self._write_lock():
+            self._ensure_index_locked(repair=True)
+        return len(self)
+
+    # Reading -----------------------------------------------------------------
+
+    def content_pins(self, sequences: Sequence[str]) -> dict[str, str]:
+        """Pin selected immutable files, excluding the rebuildable index.
+
+        Adding an unrelated committed segment does not invalidate this selection. A pinned
+        reader checks the requested rows' actual index associations, markers and part bytes.
+        """
+        payload = json.loads((self.directory / FEATURE_FILE).read_text(encoding="utf-8"))
+        recorded = StoredFeature.from_payload(payload)
+        if recorded != self.spec:
+            raise ValueError("Feature descriptor changed while pinning a selection.")
+        pins = {FEATURE_FILE: file_sha256(self.directory / FEATURE_FILE)}
+        markers = {}
+        for _, address, _ in self._located(sequences, load=False):
+            prefix = f"{SEGMENTS_DIRECTORY}/{address.segment}"
+            if address.segment not in markers:
+                markers[address.segment] = self._segment_payload(address.segment)
+                marker = self.directory / prefix / COMMIT_FILE
+                pins[f"{prefix}/{COMMIT_FILE}"] = file_sha256(marker)
+            part = markers[address.segment]["parts"][address.part]
+            pins[f"{prefix}/{PART_TEMPLATE.format(address.part)}"] = part["sha256"]
+            if part.get("row_metadata") is not None:
+                pins[f"{prefix}/{part['row_metadata']['file']}"] = part["row_metadata"]["sha256"]
+        for relative, expected in pins.items():
+            actual = file_sha256(self.directory / relative)
+            if actual != expected:
+                raise ValueError("Feature bytes changed while pinning a selection.")
+            self._check_pin(relative, actual)
+        return pins
+
+    def read(self, sequences: Sequence[str]) -> list[Tensor]:
+        """Each sequence's feature as a tensor, in the order asked.
+
+        A dense feature gives ``(w,)``, a ragged one ``(r_i, d)``, and a csr one a densified
+        ``(w,)`` row; ``read_sparse`` keeps a csr row compressed. ``ragged_topk`` requires
+        ``read_topk`` to keep residue codes sparse. A sequence the store lacks raises,
+        because a silent zero row is indistinguishable from a real one.
+        """
+
+        if self.spec.layout == CSR:
+            return [row.to_dense(self.spec.width) for row in self.read_sparse(sequences)]  # (w,) each
+        if self.spec.layout == RAGGED_TOPK:
+            raise ValueError(
+                "Use read_topk for sparse residue codes; implicit densification is refused."
+            )
+        rows: list[Tensor] = []
+        for sequence, address, tensors in self._located(sequences):
+            # tensors["values"]: (n, w) dense or (sum r_i, d) ragged, for the n rows of one part
+            if self.spec.layout == DENSE:
+                rows.append(tensors["values"][address.row].clone())  # (w,)
+            else:
+                offsets = tensors["offsets"]  # (n + 1,)
+                start, stop = int(offsets[address.row]), int(offsets[address.row + 1])
+                rows.append(tensors["values"][start:stop].clone())  # (r_i, d)
+            del sequence
+        return rows  # (w,) or (r_i, d) per sequence
+
+    def read_sparse(self, sequences: Sequence[str]) -> list[SparseRow]:
+        """Each sequence's compressed row, in the order asked. Only for a csr feature."""
+
+        if self.spec.layout != CSR:
+            raise ValueError(f"read_sparse needs a csr feature; this one is {self.spec.layout}.")
+        rows: list[SparseRow] = []
+        for _, address, tensors in self._located(sequences):
+            indptr = tensors["indptr"]  # (n + 1,)
+            start, stop = int(indptr[address.row]), int(indptr[address.row + 1])
+            positions = tensors.get("positions")  # (nnz,) or None, for the part's nnz entries
+            rows.append(
+                SparseRow(
+                    indices=tensors["indices"][start:stop].clone(),  # (nnz_i,)
+                    values=tensors["values"][start:stop].clone(),  # (nnz_i,)
+                    positions=None if positions is None else positions[start:stop].clone(),  # (nnz_i,)
+                )
+            )
+        return rows
+
+    def read_topk(self, sequences: Sequence[str]) -> list[TopKRow]:
+        """Return sparse ``(residues,k)`` codes in request order, including duplicate requests."""
+        if self.spec.layout != RAGGED_TOPK:
+            raise ValueError(
+                f"read_topk needs a ragged_topk feature; this one is {self.spec.layout}."
+            )
+        rows = []
+        for _, address, tensors in self._located(sequences):
+            offsets = tensors["offsets"]  # (n + 1,)
+            start, stop = int(offsets[address.row]), int(offsets[address.row + 1])
+            rows.append(TopKRow(
+                tensors["indices"][start:stop].clone(),  # (r_i, k)
+                tensors["values"][start:stop].clone(),  # (r_i, k)
+            ))
+        return rows
+
+    def residue_counts(self, sequences: Sequence[str]) -> list[int]:
+        """Each sequence's stored residue count, which a length-bucketed reader batches by."""
+
+        return [address.residues for _, address, _ in self._located(sequences, load=False)]
+
+    def row_metadata(
+        self, sequences: Sequence[str], *, verify_data: bool = True,
+    ) -> list[dict[str, Any]]:
+        """Read opaque row identities in request order and verify committed file digests.
+
+        The contract provider owns the scientific schema. This layer verifies the physical
+        association between sequence, index address, commit marker, data and metadata file.
+        Legacy rows without identities raise rather than becoming canonical cache hits.
+        """
+        markers: dict[str, dict[str, Any]] = {}
+        parts: dict[tuple[str, int], list[dict[str, Any]]] = {}
+        identities = []
+        for sequence, address, _ in self._located(sequences, load=False):
+            if not _NAME.fullmatch(address.segment) or address.part < 0:
+                raise ValueError("Invalid committed feature address.")
+            if address.segment not in markers:
+                marker = self.directory / SEGMENTS_DIRECTORY / address.segment / COMMIT_FILE
+                payload = json.loads(marker.read_text(encoding="utf-8"))
+                if payload["key"] != self.spec.key or payload["fingerprint"] != address.segment:
+                    raise ValueError("Committed feature identity does not match the index.")
+                markers[address.segment] = payload
+            payload = markers[address.segment]
+            matching = [part for part in payload["parts"] if part["part"] == address.part]
+            if len(matching) != 1:
+                raise ValueError("Index part does not name one committed feature part.")
+            part = matching[0]
+            if (not 0 <= address.row < len(part["digests"])
+                    or part["digests"][address.row] != sequence_digest(sequence)
+                    or part["residues"][address.row] != address.residues):
+                raise ValueError("Feature index and committed sequence row disagree.")
+            where = (address.segment, address.part)
+            if where not in parts:
+                identity = part.get("row_metadata")
+                if not isinstance(identity, dict):
+                    raise ValueError(
+                        "Stored row identities are unavailable; recompute this legacy feature."
+                    )
+                filename = f"part-{address.part:05d}.rows.json.gz"
+                if identity.get("file") != filename:
+                    raise ValueError("Invalid row metadata filename.")
+                path = self.directory / SEGMENTS_DIRECTORY / address.segment / filename
+                if file_sha256(path) != identity.get("sha256"):
+                    raise ValueError("Stored row metadata digest does not match its commit marker.")
+                if verify_data and file_sha256(self._part_path(*where)) != part.get("sha256"):
+                    raise ValueError("Stored feature data digest does not match its commit marker.")
+                rows = json.loads(gzip.decompress(path.read_bytes()))
+                if (not isinstance(rows, list) or len(rows) != len(part["digests"])
+                        or not all(isinstance(row, dict) for row in rows)):
+                    raise ValueError("Stored row metadata does not match the committed part rows.")
+                parts[where] = rows
+            identities.append(parts[where][address.row])
+        return identities
+
+    # Internals ---------------------------------------------------------------
+
+    def _located(
+        self, sequences: Sequence[str], *, load: bool = True
+    ) -> list[tuple[str, RowAddress, dict[str, Tensor]]]:
+        """Resolve each sequence to its address, reading each part file at most once."""
+
+        self._check_pin(FEATURE_FILE)
+        addresses = list(zip(sequences, self._addresses(sequences), strict=True))
+        cache: dict[tuple[str, int], dict[str, Tensor]] = {}
+        markers: dict[str, dict[str, Any]] = {}
+        located: list[tuple[str, RowAddress, dict[str, Tensor]]] = []
+        for sequence, address in addresses:
+            if address.segment not in markers:
+                markers[address.segment] = self._segment_payload(address.segment)
+            payload = markers[address.segment]
+            if not 0 <= address.part < len(payload["parts"]):
+                raise ValueError("Index part does not name one committed feature part.")
+            part = payload["parts"][address.part]
+            if (not 0 <= address.row < len(part["digests"])
+                    or part["digests"][address.row] != sequence_digest(sequence)
+                    or part["residues"][address.row] != address.residues):
+                raise ValueError("Feature index and committed sequence row disagree.")
+            where = (address.segment, address.part)
+            if where not in cache:
+                tensors = self._verified_part(address.segment, part)
+                cache[where] = tensors if load else {}
+            tensors = cache[where]
+            located.append((sequence, address, tensors))
+        return located  # (sequence, address, part tensors) per sequence; the tensors are those of _verified_part
+
+    def _addresses(self, sequences: Sequence[str]) -> list[RowAddress]:
+        """Each sequence's address in the order given, from one index connection.
+
+        Opening the index once per sequence cost minutes for a run of a hundred thousand sequences on a
+        network volume. The first sequence the store lacks raises, as ``read`` always has.
+        """
+
+        digests = [sequence_digest(sequence) for sequence in sequences]
+        found: dict[str, RowAddress] = {}
+        with self._connect() as connection:
+            for chunk in _chunks(list(dict.fromkeys(digests)), 512):
+                marks = ",".join("?" * len(chunk))
+                for digest, *where in connection.execute(
+                    f"SELECT digest, segment, part, row, residues FROM rows WHERE digest IN ({marks})",
+                    chunk,
+                ):
+                    found[digest] = RowAddress(*where)
+        addresses: list[RowAddress] = []
+        for sequence, digest in zip(sequences, digests, strict=True):
+            address = found.get(digest)
+            if address is None:
+                raise KeyError(
+                    f"Feature {self.spec.key!r} has no row for a sequence of "
+                    f"{len(sequence)} residues (sha256 {digest[:12]}). "
+                    "Embed it first, or call missing() before reading."
+                )
+            addresses.append(address)
+        return addresses
+
+    def _segment_payload(self, fingerprint: str) -> dict[str, Any]:
+        """Validate a commit marker before using any filename or row it supplies."""
+        if not isinstance(fingerprint, str) or not _NAME.fullmatch(fingerprint):
+            raise ValueError("Invalid committed segment fingerprint.")
+        marker = self.directory / SEGMENTS_DIRECTORY / fingerprint / COMMIT_FILE
+        encoded = marker.read_bytes()
+        self._check_pin(
+            f"{SEGMENTS_DIRECTORY}/{fingerprint}/{COMMIT_FILE}",
+            hashlib.sha256(encoded).hexdigest(),
+        )
+        payload = json.loads(encoded.decode("utf-8"))
+        if not isinstance(payload, dict):
+            raise ValueError("A feature commit marker must be an object.")
+        if (payload.get("format") != FORMAT or payload.get("key") != self.spec.key
+                or payload.get("fingerprint") != fingerprint):
+            raise ValueError("Committed feature identity does not match its directory.")
+        if any(key in payload for key in (
+            "transaction_schema", "manifest_sha256", "descriptor_sha256",
+        )):
+            if (type(payload.get("transaction_schema")) is not int
+                    or payload["transaction_schema"] != 2):
+                raise ValueError("Unsupported feature transaction schema.")
+            unsigned = {key: value for key, value in payload.items() if key != "manifest_sha256"}
+            if payload.get("manifest_sha256") != json_sha256(unsigned, allow_nan=False):
+                raise ValueError("Feature commit marker checksum mismatch.")
+            if payload.get("descriptor_sha256") != json_sha256(self.spec.payload(), allow_nan=False):
+                raise ValueError("Feature descriptor does not match its committed digest.")
+        parts = payload.get("parts")
+        if not isinstance(parts, list) or not parts:
+            raise ValueError("A committed segment must contain parts.")
+        seen: set[str] = set()
+        for number, part in enumerate(parts):
+            if (not isinstance(part, dict) or type(part.get("part")) is not int
+                    or part["part"] != number):
+                raise ValueError("Committed part numbers must be consecutive and unique.")
+            digests, residues = part.get("digests"), part.get("residues")
+            if (not isinstance(digests, list) or not digests or not isinstance(residues, list)
+                    or len(digests) != len(residues)):
+                raise ValueError("Committed sequence and residue counts disagree.")
+            if any(not isinstance(value, str) or not _SHA256.fullmatch(value) for value in digests):
+                raise ValueError("Invalid committed sequence digest.")
+            if len(set(digests)) != len(digests) or seen.intersection(digests):
+                raise ValueError("A committed segment repeats sequence rows.")
+            seen.update(digests)
+            if any(type(value) is not int or value < 0 for value in residues):
+                raise ValueError("Invalid committed residue count.")
+            if not isinstance(part.get("sha256"), str) or not _SHA256.fullmatch(part["sha256"]):
+                raise ValueError("Committed data checksum missing; recompute this legacy segment.")
+        if type(payload.get("rows")) is not int or payload["rows"] != len(seen):
+            raise ValueError("Committed segment row count disagrees with its parts.")
+        if (not isinstance(payload.get("committed_at"), str)
+                or not isinstance(payload.get("metadata"), dict)):
+            raise ValueError("Invalid committed segment metadata.")
+        return payload
+
+    def _verified_part(self, fingerprint: str, part: Mapping[str, Any]) -> dict[str, Tensor]:
+        """Check immutable bytes, physical layout, and optional row sidecars."""
+        number = part["part"]
+        data_digest = file_sha256(self._part_path(fingerprint, number))
+        if data_digest != part["sha256"]:
+            raise ValueError("Stored feature data digest does not match its commit marker.")
+        self._check_pin(
+            f"{SEGMENTS_DIRECTORY}/{fingerprint}/{PART_TEMPLATE.format(number)}", data_digest,
+        )
+        tensors = self._load_part(fingerprint, number)
+        # The part holds count rows. Dense: values (count, w). Otherwise offsets or indptr (count + 1,), with
+        # values (sum r_i, d) ragged, (sum r_i, k) top-k plus indices of the same shape, or (nnz,) csr.
+        spec, count = self.spec, len(part["digests"])
+        if set(tensors) != set(tensor_names(spec.layout, positions=spec.positions)):
+            raise ValueError("Committed tensor names do not match the feature layout.")
+        values = tensors["values"]  # (count, w) dense, (sum r_i, d) ragged, (sum r_i, k) top-k, (nnz,) csr
+        if values.dtype != spec.dtype:
+            raise ValueError("Committed tensor dtype does not match the feature.")
+        if spec.layout == DENSE:
+            if values.shape != (count, spec.width) or any(part["residues"]):
+                raise ValueError("Committed dense shape or residue counts disagree.")
+        else:
+            offsets = tensors["indptr" if spec.layout == CSR else "offsets"]  # (count + 1,)
+            if (offsets.dtype != torch.int64 or offsets.shape != (count + 1,)
+                    or int(offsets[0]) != 0 or int(offsets[-1]) != len(values)
+                    or bool((offsets[1:] < offsets[:-1]).any())):
+                raise ValueError("Committed row offsets are invalid.")
+            if spec.layout == RAGGED:
+                if (values.ndim != 2 or values.shape[1] != spec.width
+                        or (offsets[1:] - offsets[:-1]).tolist() != part["residues"]):
+                    raise ValueError("Committed ragged shape or residue counts disagree.")
+            elif spec.layout == RAGGED_TOPK:
+                indices = tensors["indices"]  # (sum r_i, k)
+                if (indices.dtype != torch.int32
+                        or (offsets[1:] - offsets[:-1]).tolist() != part["residues"]):
+                    raise ValueError("Committed top-k index dtype or residue counts disagree.")
+                validate_topk(indices, values, spec.width, cast(int, spec.sparse_count))
+            else:
+                indices = tensors["indices"]  # (nnz,)
+                if (values.ndim != 1 or indices.dtype != torch.int32
+                        or indices.shape != values.shape or any(part["residues"])
+                        or bool(((indices < 0) | (indices >= spec.width)).any())):
+                    raise ValueError("Committed sparse shape or indices are invalid.")
+                for start, stop in pairwise(offsets):
+                    codes = indices[int(start):int(stop)]  # (nnz_i,)
+                    if len(torch.unique(codes)) != len(codes):
+                        raise ValueError("Committed sparse row repeats indices.")
+                if spec.positions:
+                    positions = tensors["positions"]  # (nnz,)
+                    if (positions.dtype != torch.int16 or positions.shape != values.shape
+                            or bool((positions < 0).any())):
+                        raise ValueError("Committed sparse positions are invalid.")
+        if "tensor_bytes" in part and part["tensor_bytes"] != sum(
+            value.numel() * value.element_size() for value in tensors.values()
+        ):
+            raise ValueError("Committed tensor byte count disagrees with its payload.")
+        identity = part.get("row_metadata")
+        if identity is None and spec.descriptor.get("schema") in COMPLETE_SCHEMAS:
+            raise ValueError("Complete feature contracts require committed row identities.")
+        if identity is not None:
+            filename = f"part-{number:05d}.rows.json.gz"
+            if not isinstance(identity, dict) or identity.get("file") != filename:
+                raise ValueError("Invalid row metadata filename.")
+            path = self.directory / SEGMENTS_DIRECTORY / fingerprint / filename
+            metadata_digest = file_sha256(path)
+            if metadata_digest != identity.get("sha256"):
+                raise ValueError("Stored row metadata digest does not match its commit marker.")
+            self._check_pin(f"{SEGMENTS_DIRECTORY}/{fingerprint}/{filename}", metadata_digest)
+            rows = json.loads(gzip.decompress(path.read_bytes()))
+            if (not isinstance(rows, list) or len(rows) != count
+                    or not all(isinstance(row, dict) for row in rows)):
+                raise ValueError("Stored row metadata does not match the committed part rows.")
+        return tensors  # (count, w) values for dense; offsets or indptr (count + 1,) with values as above otherwise
+
+    def _check_pin(self, relative: str, actual: str | None = None) -> None:
+        """Check independent selection pins, in addition to a marker's internal checksums."""
+        if self._content_pins is not None:
+            if relative not in self._content_pins:
+                raise ValueError(
+                    f"Requested feature file is outside the pinned selection: {relative}."
+                )
+            actual = file_sha256(self.directory / relative) if actual is None else actual
+            if actual != self._content_pins[relative]:
+                raise ValueError(
+                    f"Feature file differs from its independent content pin: {relative}."
+                )
+
+    def _verified_segments(self) -> list[dict[str, Any]]:
+        payloads, seen = [], set()
+        for marker in sorted((self.directory / SEGMENTS_DIRECTORY).glob(f"*/{COMMIT_FILE}")):
+            payload = self._segment_payload(marker.parent.name)
+            for part in payload["parts"]:
+                if seen.intersection(part["digests"]):
+                    raise ValueError("Committed segments contain conflicting sequence rows.")
+                seen.update(part["digests"])
+                self._verified_part(payload["fingerprint"], part)
+            payloads.append(payload)
+        return payloads
+
+    def _write_lock(self) -> AbstractContextManager[None]:
+        self._require_writable()
+        return file_lock(self.directory / ".locks" / "store.lock")
+
+    def _segment_lock_path(self, fingerprint: str) -> Path:
+        # 128 bits keep distinct segments on distinct locks, and a short name keeps the path under
+        # Windows' 260-character limit, which a store in a nested directory would otherwise exceed.
+        name = hashlib.sha256(fingerprint.encode("utf-8")).hexdigest()[:32]
+        return self.directory / ".locks" / f"segment-{name}.lock"
+
+    def _load_part(self, segment: str, part: int) -> dict[str, Tensor]:
+        path = self._part_path(segment, part)
+        try:
+            from safetensors import safe_open
+        except ImportError as error:
+            raise ImportError("Reading a feature store requires the 'safetensors' package.") from error
+        with safe_open(path, framework="pt", device="cpu") as handle:
+            # `safe_open` is a handle with keys(), not a mapping: iterating it directly does not work.
+            return {name: cast(Tensor, handle.get_tensor(name)) for name in handle.keys()}  # noqa: dict-idiom  # (...) each, as stored
+
+    def _part_path(self, segment: str, part: int) -> Path:
+        return self.directory / SEGMENTS_DIRECTORY / segment / PART_TEMPLATE.format(part)
+
+    @contextmanager
+    def _connect(self) -> Iterator[sqlite3.Connection]:
+        """A connection that is always closed.
+
+        `sqlite3.Connection` as a context manager ends the transaction but leaves the handle open,
+        which on Windows keeps the index file locked against the next writer.
+        """
+
+        database = self.directory / INDEX_FILE
+        connection = (
+            sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)
+            if self._read_only else sqlite3.connect(database)
+        )
+        try:
+            yield connection
+        finally:
+            connection.close()
+
+    def _require_writable(self) -> None:
+        if self._read_only:
+            raise PermissionError("This feature store handle is read-only.")
+
+    def _ensure_index(self) -> None:
+        self._require_writable()
+        with self._write_lock():
+            self._ensure_index_locked()
+
+    def _ensure_index_locked(self, *, repair: bool = False) -> None:
+        if not self._deep_verify and not repair:
+            self._index_new_segments_locked()
+            return
+        payloads = self._verified_segments()
+        try:
+            with self._connect() as connection:
+                _rebuild_rows(connection, payloads, repair=repair)
+        except sqlite3.DatabaseError as error:
+            if getattr(error, "sqlite_errorcode", None) not in (
+                sqlite3.SQLITE_CORRUPT, sqlite3.SQLITE_NOTADB,
+            ):
+                raise
+            # A derived corrupt index can be replaced only after every source segment verifies.
+            temporary = self.directory / (INDEX_FILE + ".writing")
+            temporary.unlink(missing_ok=True)
+            connection = sqlite3.connect(temporary)
+            try:
+                _rebuild_rows(connection, payloads)
+            finally:
+                connection.close()
+            publish_file(temporary, self.directory / INDEX_FILE)
+
+    def _index_uncounted_segments(self) -> None:
+        """Index any committed segment the index does not hold, which a crash can leave behind."""
+
+        self._ensure_index()
+
+    def _index_new_segments_locked(self) -> None:
+        """Index only the committed segments the index lacks, reading their markers and nothing else.
+
+        This is the recovery of a store opened with ``deep_verify=False``. A segment already in the
+        index is trusted until a reader touches its parts, so the cost of a commit grows with the new
+        segment and the count of segments, never with the bytes already stored. A corrupt index is
+        rebuilt from the markers alone.
+        """
+
+        try:
+            with self._connect() as connection:
+                connection.execute("BEGIN IMMEDIATE")
+                _create_index_tables(connection)
+                self._index_missing_segments(connection)
+                connection.commit()
+        except sqlite3.DatabaseError as error:
+            if getattr(error, "sqlite_errorcode", None) not in (
+                sqlite3.SQLITE_CORRUPT, sqlite3.SQLITE_NOTADB,
+            ):
+                raise
+            temporary = self.directory / (INDEX_FILE + ".writing")
+            temporary.unlink(missing_ok=True)
+            connection = sqlite3.connect(temporary)
+            try:
+                connection.execute("BEGIN IMMEDIATE")
+                _create_index_tables(connection)
+                self._index_missing_segments(connection)
+                connection.commit()
+            finally:
+                connection.close()
+            publish_file(temporary, self.directory / INDEX_FILE)
+
+    def _index_missing_segments(self, connection: sqlite3.Connection) -> None:
+        known = {row[0] for row in connection.execute("SELECT segment FROM segments")}
+        # A partial copy's index is always this code's, and lists a segment once all its parts are in.
+        if not known and not self._partial:
+            # An index written before the segment table existed lists its segments only through its rows.
+            connection.execute("INSERT OR IGNORE INTO segments (segment) SELECT DISTINCT segment FROM rows")
+            known = {row[0] for row in connection.execute("SELECT segment FROM segments")}
+        for marker in sorted((self.directory / SEGMENTS_DIRECTORY).glob(f"*/{COMMIT_FILE}")):
+            name = marker.parent.name
+            if name in known:
+                continue
+            payload = self._segment_payload(name)
+            complete = True
+            for part in payload["parts"]:
+                number = int(part["part"])
+                if not self._part_path(name, number).is_file():
+                    if not self._partial:
+                        raise ValueError(f"Committed part {part['part']} of segment {name!r} is missing.")
+                    complete = False  # not fetched: its rows read as missing until a fetch brings it
+                    continue
+                if self._partial and connection.execute(
+                    "SELECT 1 FROM rows WHERE segment = ? AND part = ? LIMIT 1", (name, number),
+                ).fetchone():
+                    continue  # an earlier open of this copy indexed it
+                _insert_rows(connection, name, number, part)
+            if complete:
+                connection.execute("INSERT OR IGNORE INTO segments (segment) VALUES (?)", (name,))
+
+
+class SegmentWriter:
+    """One run's segment: parts as it streams, then one commit marker and the index rows."""
+
+    def __init__(
+        self, store: FeatureStore, path: Path, fingerprint: str, metadata: dict[str, Any],
+        before_commit: Callable[[], None] | None = None,
+        *, lease: object | None = None, verify_staged: bool = True,
+    ) -> None:
+        store._require_writable()
+        if lease is not _WRITER_LEASE:
+            raise RuntimeError("Open a segment writer through FeatureStore.segment().")
+        self.store = store
+        self.path = path
+        self.fingerprint = fingerprint
+        self.metadata = json.loads(json.dumps(metadata, allow_nan=False))
+        self.before_commit = before_commit
+        self.committed = False
+        self.closed = False
+        self.verify_staged = verify_staged
+        self._owner_pid = os.getpid()
+        self._failed = False
+        self._seen_digests: set[str] = set()
+        self._parts: list[dict[str, Any]] = []
+        # Parts are numbered when reserved and recorded when written, which may be on another thread.
+        self._lock = threading.Lock()
+        self._reserved = 0
+        self._staged_stats: dict[str, tuple[int, int]] = {}
+
+    def __getstate__(self) -> dict[str, Any]:
+        # A lock does not pickle. A copy sent to another process gets a fresh one and still refuses every write,
+        # commit and abandon, because its owner pid is not that process's.
+        state = dict(self.__dict__)
+        del state["_lock"]
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._lock = threading.Lock()
+
+    def _require_open(self) -> None:
+        if self._owner_pid != os.getpid():
+            raise RuntimeError("A segment writer belongs to the process that opened its context.")
+        if self.committed or (self.path / COMMIT_FILE).exists():
+            raise RuntimeError(f"Segment {self.fingerprint!r} is already committed.")
+        if self.closed or self._failed:
+            raise RuntimeError("This segment writer is closed or failed; start a fresh attempt.")
+
+    @property
+    def parts(self) -> tuple[int, ...]:
+        """The part numbers written so far, which are visible only once committed."""
+
+        return tuple(int(part["part"]) for part in self._parts)
+
+    def append_bounded(
+        self, sequences: Sequence[str],
+        rows: Sequence[Tensor] | Sequence[SparseRow] | Sequence[TopKRow],
+        *, max_tensor_bytes: int, row_metadata: Sequence[Mapping[str, Any]] | None = None,
+    ) -> tuple[int, ...]:
+        """Split a bounded window into lossless parts capped by their encoded tensor payload.
+
+        Metadata sidecars and safetensors headers are separate. A row cannot span parts; reject
+        an oversized row instead of silently writing an oversized part or changing its values.
+        """
+        # rows: (w,) dense or (r_i, d) ragged tensors; a SparseRow holds (nnz_i,) and a TopKRow (r_i, k) tensors
+        if type(max_tensor_bytes) is not int or max_tensor_bytes <= 0:
+            raise ValueError("max_tensor_bytes must be a positive integer.")
+        if not sequences or len(sequences) != len(rows):
+            raise ValueError("append_bounded needs one row per sequence and at least one row.")
+        if len(set(sequences)) != len(sequences):
+            raise ValueError("This batch repeats sequences; each feature has one row.")
+        if row_metadata is not None and len(row_metadata) != len(sequences):
+            raise ValueError("Row metadata must contain one identity per sequence.")
+        spec = self.store.spec
+        initial = 0 if spec.layout == DENSE else 8
+        sizes = [
+            row_tensor_bytes(row, spec.layout, spec.width, spec.dtype, positions=spec.positions)
+            for row in rows
+        ]
+        if any(initial + size > max_tensor_bytes for size in sizes):
+            raise ValueError(
+                "A feature row exceeds max_tensor_bytes; increase the explicit part budget."
+            )
+        written, start, used = [], 0, initial
+        for stop in range(len(rows) + 1):
+            size = sizes[stop] if stop < len(rows) else 0
+            if stop == len(rows) or used + size > max_tensor_bytes:
+                part = self.append(
+                    sequences[start:stop], rows[start:stop],
+                    row_metadata=None if row_metadata is None else row_metadata[start:stop],
+                )
+                if self._parts[part]["tensor_bytes"] > max_tensor_bytes:
+                    raise RuntimeError("Encoded tensor payload exceeded the planned part budget.")
+                written.append(part)
+                start, used = stop, initial
+            used += size
+        return tuple(written)
+
+    def append(
+        self, sequences: Sequence[str],
+        rows: Sequence[Tensor] | Sequence[SparseRow] | Sequence[TopKRow],
+        *, row_metadata: Sequence[Mapping[str, Any]] | None = None,
+    ) -> int:
+        """Write one part holding these rows, and return the part number.
+
+        Sequences the store already holds, or that repeat inside this batch, are an error: a
+        feature has one row, and writing it twice makes two answers to one question. Call
+        ``missing`` first.
+        """
+
+        # rows: (w,) dense or (r_i, d) ragged tensors; a SparseRow holds (nnz_i,) and a TopKRow (r_i, k) tensors
+        self._require_open()
+        if len(sequences) != len(rows):
+            raise ValueError(
+                f"append needs one row per sequence; received {len(sequences)} sequences and "
+                f"{len(rows)} rows."
+            )
+        if not sequences:
+            raise ValueError("append needs at least one sequence.")
+        digests = [sequence_digest(sequence) for sequence in sequences]
+        repeated = sorted({digest for digest in digests if digests.count(digest) > 1})
+        if repeated or self._seen_digests.intersection(digests):
+            raise ValueError(f"This batch repeats {len(repeated)} sequences; each feature has one row.")
+        already = self.store.missing(sequences)
+        if len(already) != len(sequences):
+            raise ValueError(
+                f"{len(sequences) - len(already)} of these sequences already have a row in "
+                f"{self.store.spec.key!r}; call missing() and embed only what it returns."
+            )
+
+        spec = self.store.spec
+        if spec.descriptor.get("schema") in COMPLETE_SCHEMAS and row_metadata is None:
+            raise ValueError(
+                "Complete feature contracts require a persisted identity for each row."
+            )
+        if row_metadata is not None and len(row_metadata) != len(sequences):
+            raise ValueError("Row metadata must contain one identity per sequence.")
+        encoded_metadata = None if row_metadata is None else gzip.compress(
+            json.dumps([dict(row) for row in row_metadata], sort_keys=True, allow_nan=False,
+                       separators=(",", ":")).encode("utf-8"), mtime=0,
+        )
+        if spec.layout == DENSE:
+            tensors = encode_dense(cast(Sequence[Tensor], rows), spec.width, spec.dtype)
+            residues = [0] * len(sequences)
+        elif spec.layout == CSR:
+            tensors = encode_csr(cast(Sequence[SparseRow], rows), spec.width, spec.dtype)
+            if spec.positions and "positions" not in tensors:
+                raise ValueError(
+                    f"Feature {spec.key!r} stores argmax positions; these rows carry none."
+                )
+            if not spec.positions and "positions" in tensors:
+                raise ValueError(
+                    f"Feature {spec.key!r} stores no argmax positions; these rows carry them."
+                )
+            residues = [0] * len(sequences)
+        else:
+            if spec.layout == RAGGED_TOPK:
+                tensors = encode_topk_rows(
+                    cast(Sequence[TopKRow], rows), spec.width,
+                    cast(int, spec.sparse_count), spec.dtype,
+                )
+            else:
+                tensors = encode_ragged(cast(Sequence[Tensor], rows), spec.width, spec.dtype)
+            offsets = tensors["offsets"]  # (n + 1,)
+            residues = [
+                int(offsets[position + 1]) - int(offsets[position])
+                for position in range(len(sequences))
+            ]
+        written = row_count(spec.layout, tensors)
+        if written != len(sequences):
+            raise ValueError(
+                f"Encoded {written} rows for {len(sequences)} sequences; the layout and the rows "
+                "disagree."
+            )
+
+        part = self.reserve_part()
+        try:
+            _transaction_event("before_part_write", self.path)
+            _save_safetensors_atomically(
+                self.path / PART_TEMPLATE.format(part),
+                {name: tensors[name]
+                 for name in tensor_names(spec.layout, positions=spec.positions)},
+            )
+            _transaction_event("after_part_write", self.path)
+            record = {
+                "part": part, "digests": digests, "residues": residues,
+                "sha256": file_sha256(self.path / PART_TEMPLATE.format(part)),
+                "tensor_bytes": sum(
+                    value.numel() * value.element_size() for value in tensors.values()),
+            }
+            if encoded_metadata is not None:
+                target = self.path / f"part-{part:05d}.rows.json.gz"
+                temporary = target.with_name(target.name + ".writing")
+                temporary.write_bytes(encoded_metadata)
+                publish_file(temporary, target)
+                record["row_metadata"] = {"file": target.name, "sha256": file_sha256(target)}
+                _transaction_event("after_metadata_write", self.path)
+        except BaseException:
+            self._failed = True
+            raise
+        self._parts.append(record)
+        self._seen_digests.update(digests)
+        return part
+
+    def reserve_part(self) -> int:
+        """The next part number, taken before a part is written so that writers on several threads never collide."""
+
+        with self._lock:
+            number = self._reserved
+            self._reserved += 1
+        return number
+
+    def seen(self, digests: Sequence[str]) -> bool:
+        """Whether any digest was already written by this segment, a cheap guard before a large write."""
+
+        with self._lock:
+            return not self._seen_digests.isdisjoint(digests)
+
+    def append_packed(
+        self, part: int, digests: Sequence[str], tensors: Mapping[str, Tensor], residues: Sequence[int],
+        *, row_metadata: Sequence[Mapping[str, Any]] | None = None,
+    ) -> dict[str, Any]:
+        """Write one large part from tensors already packed in the layout, hashing as it writes.
+
+        ``tensors`` holds exactly the layout's tensors, as ``encode_*`` would build them. This skips
+        the per-row copies, finite scans and index re-reads of ``append``: the caller proved the values
+        finite on the device and packed them in order. Shapes, dtypes, offsets and the residue counts
+        are still checked, because they decide whether a reader can address the part. The file and
+        its digest come from one pass, and one flush covers the part. Safe to call from several
+        threads with parts from ``reserve_part``.
+
+        Shapes, with b rows (sequences) in the part and n = sum(residues) stored rows: dense ``values`` (b, w);
+        ragged ``offsets`` (b + 1,) and ``values`` (n, w); ragged top-k adds ``indices`` (n, k); csr ``indptr``
+        (b + 1,) with ``indices`` and ``values`` (nnz,). ``residues[i]`` is row i's stored rows, which is l + 2
+        for a stream that keeps CLS and EOS.
+        """
+        # tensors: (b, w) dense; (b + 1,) offsets and (n, w) values ragged; (b + 1,) indptr and (nnz,) csr.
+        spec = self.store.spec
+        if self._owner_pid != os.getpid():
+            raise RuntimeError("A segment writer belongs to the process that opened its context.")
+        if self.committed or (self.path / COMMIT_FILE).exists():
+            raise RuntimeError(f"Segment {self.fingerprint!r} is already committed.")
+        if self.closed or self._failed:
+            raise RuntimeError("This segment writer is closed or failed; start a fresh attempt.")
+        count = len(digests)
+        if not count or len(residues) != count:
+            raise ValueError("append_packed needs one residue count and one digest per row, and at least one row.")
+        if spec.descriptor.get("schema") in COMPLETE_SCHEMAS and row_metadata is None:
+            raise ValueError("Complete feature contracts require a persisted identity for each row.")
+        if row_metadata is not None and len(row_metadata) != count:
+            raise ValueError("Row metadata must contain one identity per sequence.")
+        if set(tensors) != set(tensor_names(spec.layout, positions=spec.positions)):
+            raise ValueError("Packed tensors do not match the feature layout.")
+        if len(set(digests)) != count or self.seen(digests):
+            raise ValueError("A packed part repeats sequences; each feature has one row.")
+        _check_packed(spec, tensors, count, residues)
+
+        try:
+            _transaction_event("before_part_write", self.path)
+            target = self.path / PART_TEMPLATE.format(part)
+            digest, size = _write_safetensors_streaming(target, tensors)
+            _transaction_event("after_part_write", self.path)
+            record: dict[str, Any] = {
+                "part": part, "digests": list(digests), "residues": [int(value) for value in residues],
+                "sha256": digest,
+                "tensor_bytes": sum(value.numel() * value.element_size() for value in tensors.values()),
+            }
+            stats = {target.name: (size, target.stat().st_mtime_ns)}
+            if row_metadata is not None:
+                encoded = gzip.compress(
+                    json.dumps([dict(row) for row in row_metadata], sort_keys=True, allow_nan=False,
+                               separators=(",", ":")).encode("utf-8"), mtime=0,
+                )
+                sidecar = self.path / f"part-{part:05d}.rows.json.gz"
+                temporary = sidecar.with_name(sidecar.name + ".writing")
+                temporary.write_bytes(encoded)
+                publish_file(temporary, sidecar, sync_parent=False)
+                record["row_metadata"] = {"file": sidecar.name, "sha256": hashlib.sha256(encoded).hexdigest()}
+                stats[sidecar.name] = (len(encoded), sidecar.stat().st_mtime_ns)
+                _transaction_event("after_metadata_write", self.path)
+        except BaseException:
+            self._failed = True
+            raise
+        with self._lock:
+            self._parts.append(record)
+            self._seen_digests.update(digests)
+            self._staged_stats.update(stats)
+        return record
+
+    def commit(self) -> SegmentReceipt:
+        """Write the commit marker, then index the parts. Nothing reads a segment until this."""
+
+        self._require_open()
+        with self._lock:
+            # Parts written on several threads finish in any order; the marker lists them by number.
+            self._parts.sort(key=lambda part: part["part"])
+        if not self._parts:
+            raise RuntimeError(f"Segment {self.fingerprint!r} holds no parts to commit.")
+        payload = {
+            "format": FORMAT,
+            "key": self.store.spec.key,
+            "fingerprint": self.fingerprint,
+            "committed_at": datetime.now(UTC).isoformat(timespec="seconds"),
+            "rows": sum(len(part["digests"]) for part in self._parts),
+            "metadata": self.metadata,
+            "parts": self._parts,
+            "transaction_schema": 2,
+            "descriptor_sha256": json_sha256(self.store.spec.payload(), allow_nan=False),
+        }
+        recorded = json.loads((self.store.directory / FEATURE_FILE).read_text(encoding="utf-8"))
+        if recorded != self.store.spec.payload():
+            raise ValueError("Feature descriptor changed while the segment was staged.")
+        for part in self._parts:
+            self._check_staged(self.path / PART_TEMPLATE.format(part["part"]), part["sha256"])
+            identity = part.get("row_metadata")
+            if identity:
+                self._check_staged(self.path / identity["file"], identity["sha256"])
+        if self.before_commit is not None:
+            self.before_commit()
+        # Parts and sidecars were renamed without a directory flush; one flush covers them all.
+        sync_directory(self.path)
+        payload["manifest_sha256"] = json_sha256(payload, allow_nan=False)
+        with self.store._write_lock():
+            # Recover prior durable markers before deciding whether any row conflicts.
+            self.store._ensure_index_locked()
+            with self.store._connect() as connection:
+                connection.execute("BEGIN IMMEDIATE")
+                _transaction_event("before_index_update", self.path)
+                _create_index_tables(connection)  # idempotent: an index from before the segment list gains it here
+                for part in self._parts:
+                    _insert_rows(connection, self.fingerprint, int(part["part"]), part)
+                connection.execute("INSERT OR IGNORE INTO segments (segment) VALUES (?)", (self.fingerprint,))
+                _transaction_event("after_index_update", self.path)
+                _transaction_event("before_commit_marker", self.path)
+                _write_json_atomically(self.path / COMMIT_FILE, payload)
+                self.committed = True
+                _transaction_event("after_commit_marker", self.path)
+                connection.commit()
+                _transaction_event("after_index_commit", self.path)
+        return _receipt(payload)
+
+    def _check_staged(self, path: Path, digest: str) -> None:
+        """Refuse a staged file that changed since it was written.
+
+        A writer that hashed as it wrote (``verify_staged=False``) compares size and modification
+        time, which catches an appended or rewritten file without a second read of the payload.
+        """
+        recorded = self._staged_stats.get(path.name)
+        if self.verify_staged or recorded is None:
+            if file_sha256(path) != digest:
+                raise ValueError(
+                    "Staged feature data changed before commit." if path.suffix == ".safetensors"
+                    else "Staged row metadata changed before commit."
+                )
+            return
+        observed = path.stat()
+        if (observed.st_size, observed.st_mtime_ns) != recorded:
+            raise ValueError(
+                "Staged feature data changed before commit." if path.suffix == ".safetensors"
+                else "Staged row metadata changed before commit."
+            )
+
+    def abandon(self) -> None:
+        """Delete this segment's parts, for a run that decides not to keep them."""
+
+        if self._owner_pid != os.getpid():
+            raise RuntimeError("A segment writer belongs to the process that opened its context.")
+        if self.committed or (self.path / COMMIT_FILE).exists():
+            raise RuntimeError(f"Segment {self.fingerprint!r} is committed and immutable.")
+        if self.closed:
+            raise RuntimeError("This segment writer is closed.")
+        _discard_segment(self.path)
+        self._parts.clear()
+        self._seen_digests.clear()
+        self.closed = True
+
+
+# The documented public entry point (docs/feature_store.md, `features.__all__`), kept under its name.
+def open_feature(root: str | Path, spec: StoredFeature) -> FeatureStore:  # noqa: renaming-wrapper
+    """Open, or create, the store for one feature under ``root``."""
+
+    return FeatureStore.open(root, spec)
+
+
+def features_in(root: str | Path) -> tuple[FeatureStore, ...]:
+    """Every feature store under ``root``, by key."""
+
+    return tuple(
+        FeatureStore.read_only(recorded.parent)
+        for recorded in sorted(Path(root).glob(f"*/{FEATURE_FILE}"))
+    )
+
+
+def _insert_rows(
+    connection: sqlite3.Connection, fingerprint: str, part: int, payload: Mapping[str, Any]
+) -> None:
+    try:
+        connection.executemany(
+            "INSERT INTO rows (digest, segment, part, row, residues) VALUES (?, ?, ?, ?, ?)",
+            [(digest, fingerprint, part, row, int(residues))
+             for row, (digest, residues) in enumerate(
+                 zip(payload["digests"], payload["residues"], strict=True)
+             )],
+        )
+    except sqlite3.IntegrityError as error:
+        raise ValueError(
+            f"Segment {fingerprint!r} conflicts with an existing sequence row."
+        ) from error
+
+
+def _receipt(payload: Mapping[str, Any]) -> SegmentReceipt:
+    return SegmentReceipt(
+        fingerprint=str(payload["fingerprint"]),
+        parts=tuple(int(part["part"]) for part in payload["parts"]),
+        rows=int(payload["rows"]),
+        committed_at=str(payload["committed_at"]),
+        metadata=payload.get("metadata") or {},
+    )
+
+
+def _write_json_atomically(path: Path, payload: Mapping[str, Any]) -> None:
+    temporary = path.with_name(f"{path.name}.writing")
+    temporary.write_text(indented_json(payload, allow_nan=False), encoding="utf-8")
+    publish_file(temporary, path)
+
+
+def _save_safetensors_atomically(path: Path, tensors: Mapping[str, Tensor]) -> None:
+    # tensors: (n, w), (n + 1,), (nnz,) or (sum r_i, d) by name and layout; each is written as it is
+    try:
+        from safetensors.torch import save_file
+    except ImportError as error:
+        raise ImportError("Writing a feature store requires the 'safetensors' package.") from error
+    temporary = path.with_name(f"{path.name}.writing")
+    save_file({name: tensor.contiguous() for name, tensor in tensors.items()}, str(temporary))
+    publish_file(temporary, path)
+
+
+_SAFETENSORS_DTYPES = {
+    torch.float64: "F64", torch.float32: "F32", torch.float16: "F16", torch.bfloat16: "BF16",
+    torch.int64: "I64", torch.int32: "I32", torch.int16: "I16",
+}
+_WRITE_BLOCK_BYTES = 64 * 1024**2
+
+
+def _write_safetensors_streaming(path: Path, tensors: Mapping[str, Tensor]) -> tuple[str, int]:
+    """Write ``tensors`` as a safetensors file in one pass, hashing the bytes as they are written.
+
+    The file is the safetensors format that ``safe_open`` reads: an 8-byte little-endian header
+    length, a JSON header padded with spaces to a multiple of 8, then the raw tensor bytes. Tensors go
+    in decreasing element size, so every tensor starts aligned to its own element size. The tensor
+    memory is written from its own buffer in blocks, with no second serialization copy, and the
+    SHA-256 of the file comes from those same blocks. One flush makes the part durable. Returns the
+    digest and the size in bytes.
+    """
+    # tensors: (b, w), (n, w), (n, k), (b + 1,) or (nnz,) by layout; any shape, the header records each as is.
+    ordered = sorted(tensors.items(), key=lambda item: (-item[1].element_size(), item[0]))
+    header: dict[str, Any] = {}
+    cursor = 0
+    for name, tensor in ordered:
+        size = tensor.numel() * tensor.element_size()
+        header[name] = {
+            "dtype": _SAFETENSORS_DTYPES[tensor.dtype], "shape": list(tensor.shape),
+            "data_offsets": [cursor, cursor + size],
+        }
+        cursor += size
+    encoded = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    encoded += b" " * (-len(encoded) % 8)
+    prefix = len(encoded).to_bytes(8, "little")
+    temporary = path.with_name(f"{path.name}.writing")
+    digest = hashlib.sha256()
+    with temporary.open("wb") as handle:
+        for block in (prefix, encoded):
+            handle.write(block)
+            digest.update(block)
+        for _, tensor in ordered:
+            raw = tensor.detach().contiguous().reshape(-1).view(torch.uint8)  # (bytes,) over the tensor's own memory
+            buffer = memoryview(raw.numpy())
+            for start in range(0, len(buffer), _WRITE_BLOCK_BYTES):
+                block = buffer[start : start + _WRITE_BLOCK_BYTES]
+                handle.write(block)
+                digest.update(block)
+        handle.flush()
+        flush_and_evict(handle)
+    temporary.replace(path)
+    return digest.hexdigest(), 8 + len(encoded) + cursor
+
+
+def _check_packed(
+    spec: StoredFeature, tensors: Mapping[str, Tensor], count: int, residues: Sequence[int],
+) -> None:
+    """The structural checks `append_packed` keeps: whatever decides whether a reader can address the part."""
+    # tensors: (b, w) dense; (b + 1,) offsets and (n, w) values ragged; (b + 1,) indptr and (nnz,) csr.
+    values = tensors["values"]  # (b, w) dense, (n, w) ragged, (nnz,) csr
+    if values.dtype != spec.dtype:
+        raise ValueError("Packed tensor dtype does not match the feature.")
+    if spec.layout == DENSE:
+        if tuple(values.shape) != (count, spec.width) or any(residues):
+            raise ValueError("Packed dense shape or residue counts disagree.")
+        return
+    offsets = tensors["indptr" if spec.layout == CSR else "offsets"]  # (b + 1,) int64 row boundaries
+    if (offsets.dtype != torch.int64 or tuple(offsets.shape) != (count + 1,)
+            or int(offsets[0]) != 0 or int(offsets[-1]) != len(values)
+            or bool((offsets[1:] < offsets[:-1]).any())):
+        raise ValueError("Packed row offsets are invalid.")
+    spans = (offsets[1:] - offsets[:-1]).tolist()  # (b,) stored rows (ragged) or entries (csr) per row
+    if spec.layout == RAGGED:
+        if values.ndim != 2 or values.shape[1] != spec.width or spans != list(residues):
+            raise ValueError("Packed ragged shape or residue counts disagree.")
+    elif spec.layout == RAGGED_TOPK:
+        indices = tensors["indices"]  # (n, k) int32 codes, the shape of values
+        if (indices.dtype != torch.int32 or tuple(indices.shape) != tuple(values.shape)
+                or values.ndim != 2 or values.shape[1] != spec.sparse_count or spans != list(residues)):
+            raise ValueError("Packed top-k shape, index dtype or residue counts disagree.")
+    else:
+        indices = tensors["indices"]
+        if (values.ndim != 1 or indices.dtype != torch.int32 or indices.shape != values.shape
+                or any(residues)):
+            raise ValueError("Packed sparse shape or residue counts disagree.")
+        if spec.positions and (tensors["positions"].dtype != torch.int16
+                               or tensors["positions"].shape != values.shape):
+            raise ValueError("Packed sparse positions are invalid.")
+
+
+def _discard_segment(path: Path) -> None:
+    if path.is_symlink() or any(not part.is_file() or part.is_symlink() for part in path.iterdir()):
+        raise ValueError("Refusing to discard an uncommitted segment with unexpected entries.")
+    for part in path.iterdir():
+        part.unlink()
+    path.rmdir()
+    sync_directory(path.parent)
+
+
+def _rebuild_rows(
+    connection: sqlite3.Connection, payloads: Sequence[Mapping[str, Any]], *, repair: bool = False,
+) -> None:
+    connection.execute("BEGIN IMMEDIATE")
+    _create_index_tables(connection)
+    expected = {
+        digest: (payload["fingerprint"], part["part"], row, residues)
+        for payload in payloads for part in payload["parts"]
+        for row, (digest, residues) in enumerate(
+            zip(part["digests"], part["residues"], strict=True))
+    }
+    if repair:
+        connection.execute("DELETE FROM rows")
+    indexed = {
+        row[0]: tuple(row[1:])
+        for row in connection.execute("SELECT digest, segment, part, row, residues FROM rows")
+    }
+    if any(expected.get(digest) != address for digest, address in indexed.items()):
+        raise ValueError(
+            "Feature index disagrees with committed rows; inspect it and call reindex()."
+        )
+    # Recover missing rows only. Changed addresses must fail, not silently become cache hits.
+    connection.executemany(
+        "INSERT INTO rows (digest, segment, part, row, residues) VALUES (?, ?, ?, ?, ?)",
+        [(digest, *address) for digest, address in expected.items() if digest not in indexed],
+    )
+    connection.executemany(
+        "INSERT OR IGNORE INTO segments (segment) VALUES (?)", [(payload["fingerprint"],) for payload in payloads],
+    )
+    connection.commit()
+
+
+def _create_index_tables(connection: sqlite3.Connection) -> None:
+    """The row index, and the list of segments it already holds, so recovery can skip them."""
+
+    connection.execute(
+        "CREATE TABLE IF NOT EXISTS rows (digest TEXT PRIMARY KEY, segment TEXT NOT NULL, "
+        "part INTEGER NOT NULL, row INTEGER NOT NULL, residues INTEGER NOT NULL)"
+    )
+    connection.execute("CREATE INDEX IF NOT EXISTS rows_by_segment ON rows (segment, part)")
+    connection.execute("CREATE TABLE IF NOT EXISTS segments (segment TEXT PRIMARY KEY)")
+
+
+def _transaction_event(stage: str, path: Path) -> None:
+    """A test observation point at a real I/O boundary; production performs no action."""
+
+
+def _chunks(values: Sequence[str], size: int) -> Iterator[list[str]]:
+    for start in range(0, len(values), size):
+        yield list(values[start : start + size])
+
+
+__all__ = [
+    "COMMIT_FILE",
+    "FEATURE_FILE",
+    "FORMAT",
+    "INDEX_FILE",
+    "SEGMENTS_DIRECTORY",
+    "FeatureStore",
+    "RowAddress",
+    "SegmentReceipt",
+    "SegmentWriter",
+    "StoredFeature",
+    "features_in",
+    "open_feature",
+    "partition_sequences",
+    "sequence_digest",
+]

fastplms/features/transactions.py

@@ -0,0 +1,100 @@
+"""Process ownership and durable publication for the local feature store.
+
+Locks are kernel-owned, not PID files, and release when a process exits. Lock files stay
+in place: unlinking one would let another process lock a different inode at the same path.
+See https://docs.python.org/3/library/fcntl.html and
+https://www.sqlite.org/atomiccommit.html for the locking and flush assumptions.
+"""
+
+from __future__ import annotations
+
+import errno
+import os
+import time
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import BinaryIO
+
+
+@contextmanager
+def file_lock(path: Path, *, wait: bool = True) -> Iterator[None]:
+    """Hold a process lock on a stable file, optionally refusing a competing owner."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    owner_pid = os.getpid()
+    with path.open("a+b") as handle:
+        if os.name == "nt":
+            import msvcrt
+
+            while True:
+                try:
+                    handle.seek(0)
+                    msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)
+                    break
+                except OSError as error:
+                    if error.errno not in (errno.EACCES, errno.EAGAIN, errno.EDEADLK):
+                        raise
+                    if not wait:
+                        raise BlockingIOError(
+                            "Feature segment already has an active writer."
+                        ) from error
+                    time.sleep(0.05)
+            try:
+                yield
+            finally:
+                if os.getpid() == owner_pid:
+                    handle.seek(0)
+                    msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1)
+        else:
+            import fcntl
+
+            flags = fcntl.LOCK_EX | (0 if wait else fcntl.LOCK_NB)
+            try:
+                fcntl.flock(handle.fileno(), flags)
+            except BlockingIOError as error:
+                raise BlockingIOError("Feature segment already has an active writer.") from error
+            try:
+                yield
+            finally:
+                if os.getpid() == owner_pid:
+                    fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+
+
+def sync_directory(path: Path) -> None:
+    """Flush directory entries on POSIX; Windows has no equivalent directory fsync."""
+    if os.name != "nt":
+        descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
+        try:
+            os.fsync(descriptor)
+        finally:
+            os.close(descriptor)
+
+
+def flush_and_evict(handle: BinaryIO) -> None:
+    """Make a written file durable, then drop its pages from the page cache.
+
+    A feature store writes far more than it reads back soon, and written pages stay cached. On a GH200 the
+    kernel fills the GPU's HBM, which Linux exposes as a NUMA node, with that cache, so ``nvidia-smi`` reads
+    the device as full and the next large allocation waits while the kernel evicts. The pages are clean
+    after ``fsync``, so ``POSIX_FADV_DONTNEED`` drops all of them. Windows has no ``posix_fadvise`` and
+    keeps only the flush. Callers flush Python buffers first.
+    """
+    descriptor = handle.fileno()
+    os.fsync(descriptor)
+    if hasattr(os, "posix_fadvise"):
+        os.posix_fadvise(descriptor, 0, 0, os.POSIX_FADV_DONTNEED)  # offset 0, length 0: the whole file
+
+
+def publish_file(temporary: Path, destination: Path, *, sync_parent: bool = True) -> None:
+    """Flush a complete staged file, rename it, then flush its containing directory.
+
+    A caller that publishes many files into one directory passes ``sync_parent=False`` and flushes the
+    directory once, before the commit marker, so a part costs one flush and not two. The staged file
+    leaves the page cache once it is durable (``flush_and_evict``).
+    """
+    with temporary.open("r+b") as handle:
+        flush_and_evict(handle)
+    temporary.replace(destination)
+    if sync_parent:
+        sync_directory(destination.parent)

fastplms/features/writing.py

@@ -0,0 +1,52 @@
+"""Commit a window of rows to a feature store as one segment.
+
+A pipeline that embeds a window of sequences at a time, and wants a run that dies to lose at most
+that window, does the same thing after every window: keep the sequences the store lacks, write them
+as a segment, and name the segment so a restart cannot commit the same rows under two names.
+``write_rows`` is that step. The segment is named by the digest of its sequences, so the name
+depends only on what the segment holds, and a window a restart repeats is skipped by calling
+``FeatureStore.missing`` first, exactly as a run does before it embeds.
+"""
+
+from __future__ import annotations
+
+import hashlib
+
+from collections.abc import Mapping
+from typing import Any
+
+from .conversion import Row
+from .store import FeatureStore, sequence_digest
+
+
+DEFAULT_MAX_TENSOR_BYTES = 256 * 1024**2
+
+
+def write_rows(
+    store: FeatureStore,
+    rows: Mapping[str, Row],
+    *,
+    metadata: Mapping[str, Any] | None = None,
+    max_tensor_bytes: int = DEFAULT_MAX_TENSOR_BYTES,
+) -> int:
+    """Commit ``rows`` (sequence to row) as one segment and return how many rows it holds.
+
+    A tensor for a dense or ragged feature, a ``SparseRow`` for csr, and a ``TopKRow`` for ragged
+    top-k, as ``SegmentWriter.append`` takes them, in the order the mapping yields. An empty mapping
+    writes nothing and returns zero. A sequence the store already holds raises, because a feature
+    has one row per sequence; call ``store.missing`` first.
+
+    ``metadata`` is plain data recorded in the segment's commit marker, for what the run wants to
+    remember about these rows. A window too large for ``max_tensor_bytes`` is split into parts of a
+    segment that commits or vanishes as one.
+    """
+
+    if not rows:
+        return 0
+    sequences = list(rows)
+    digest = hashlib.sha256("\n".join(sequence_digest(sequence) for sequence in sequences).encode("utf-8"))
+    with store.segment("rows-" + digest.hexdigest()[:16], metadata) as writer:
+        writer.append_bounded(
+            sequences, [rows[sequence] for sequence in sequences], max_tensor_bytes=max_tensor_bytes,
+        )
+    return len(sequences)

fastplms/json_files.py

@@ -0,0 +1,45 @@
+"""The two JSON text forms FastPLMs writes: compact for hashing, indented for files people read.
+
+This file exists twice, byte for byte: here and as ``features/json_files.py``. ``features`` loads as a
+standalone package (``foundry.embedding.private_store``) and imports nothing outside itself, so it cannot
+reach this module. ``tests/tier1_unit/test_features_package_is_self_contained.py`` fails when the two differ.
+"""
+
+from __future__ import annotations
+
+import json
+
+from typing import Any
+
+
+def compact_json(value: Any, *, ensure_ascii: bool = True, allow_nan: bool = True) -> str:
+    """Serialize with sorted keys and no whitespace, the form a digest or an identity is taken over."""
+
+    return json.dumps(
+        value,
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=ensure_ascii,
+        allow_nan=allow_nan,
+    )
+
+
+def indented_json(
+    value: Any,
+    *,
+    ensure_ascii: bool = True,
+    allow_nan: bool = True,
+    sort_keys: bool = True,
+) -> str:
+    """Serialize with two-space indentation and one trailing newline, the form of a stored JSON file."""
+
+    return (
+        json.dumps(
+            value,
+            indent=2,
+            sort_keys=sort_keys,
+            ensure_ascii=ensure_ascii,
+            allow_nan=allow_nan,
+        )
+        + "\n"
+    )

fastplms/models.toml

@@ -4,11 +4,58 @@ legal_files = [
   "THIRD_PARTY_NOTICES.md=sha256:d35e506b728868b52290672d54a89c6aade09529af99dbc8a4db72d9e9ca3460",
 ]
 
+# Released hidden-state SAEs used by the canonical suite. The training-backbone revision
+# is not reported upstream; base_model names the selected inference base, not that history.
+[[sparse_autoencoders]]
+id = "esmc_small_sae_depth"
+base_model = "esmc_small"
+layer = 23
+input_width = 960
+k = 64
+codebook_dim = 16384
+input_kind = "hidden_state"
+checkpoint_repo = "biohub/ESMC-300M-sae-layer23-k64-codebook16384"
+checkpoint_revision = "71b054fc726a9f11153a601118f03657ac6e6e80"
+checkpoint_files = [
+  "config.json=sha256:e2864c5f9756c052f343310bafeddafc3edc9514ad85981a51a1a3411186434d",
+  "layer_23.safetensors=sha256:825257f785600a7a9d462882d5cf3431ebbbba94761434e33b0c4814a7cc6270",
+]
+
+[[sparse_autoencoders]]
+id = "esmc_large_sae_depth"
+base_model = "esmc_large"
+layer = 27
+input_width = 1152
+k = 64
+codebook_dim = 16384
+input_kind = "hidden_state"
+checkpoint_repo = "biohub/ESMC-600M-sae-layer27-k64-codebook16384"
+checkpoint_revision = "b96480999aca49b7684e093f1adfcb09b55a1720"
+checkpoint_files = [
+  "config.json=sha256:0e74c8af3ff4ac644f26be308e0fe6da6c75a4a52f0d3f05b0cbbdda27efc9ce",
+  "layer_27.safetensors=sha256:5e93c230aa200876cb300377c8e22a141a7e2e316c6a3b962b2a78c4518535db",
+]
+
+[[sparse_autoencoders]]
+id = "esmc_6b_sae_depth"
+base_model = "esmc_6b"
+layer = 60
+input_width = 2560
+k = 64
+codebook_dim = 16384
+input_kind = "hidden_state"
+checkpoint_repo = "biohub/ESMC-6B-sae-layer60-k64-codebook16384"
+checkpoint_revision = "99752fe6e4d25fbb26f887db2d9225e7a577da73"
+checkpoint_files = [
+  "config.json=sha256:42567d3f757bea5c6b618abf55ed905ce4304d1eed891a2c1feae8e1ef393fd0",
+  "layer_60.safetensors=sha256:ddc1417b42cffe2d7fc2ec31783a04ac88800e6a8f8c8542f1f92d8f4f16094c",
+]
+
 [[attention_kernels]]
 implementation = "flash_attention_2"
 repository = "kernels-community/flash-attn2"
-revision = "db6b51744f0cd7061386442c09df890fc6d9f47e"
-version = 2
+revision = "81fb77c12b2ad5d69380669b46739d5868614502"
+version = 3
 expected_variant = "flash_attn2"
 dtypes = ["bfloat16"]
 min_cuda_capability = [8, 0]
@@ -64,6 +111,9 @@ distribution_files = [
 id = "biohub-transformers"
 path = "vendor/upstream/biohub-transformers"
 url = "https://github.com/Biohub/transformers.git"
+# Biohub/transformers left GitHub by 2026-10-08. Its pinned commit stays reachable through the
+# fork network of huggingface/transformers, which serves the clone and the source archive.
+fetch_url = "https://github.com/huggingface/transformers.git"
 revision = "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf"
 license = "Apache-2.0"
 license_files = ["LICENSE"]
@@ -177,7 +227,7 @@ conversion_provenance = "Input: the pinned official ESM2 state dictionary. Trans
 representative = "esm2_8m"
 documentation = "docs/models.md#esm2"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/esm2", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/_esm_rotary.py", "models/esm2", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.esm2.modeling_fastesm.FastEsmConfig", AutoModel = "fastplms.models.esm2.modeling_fastesm.FastEsmModel", AutoModelForMaskedLM = "fastplms.models.esm2.modeling_fastesm.FastEsmForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.esm2.modeling_fastesm.FastEsmForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm2.modeling_fastesm.FastEsmForTokenClassification" }
 
 [families.esm_plusplus]
@@ -204,7 +254,7 @@ conversion_provenance = "Input: the pinned Biohub ESMC checkpoint. Transformatio
 representative = "esmc_small"
 documentation = "docs/models.md#esm-and-esmc"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esm_plusplus", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/esm_plusplus", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusConfig", AutoModel = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusModel", AutoModelForMaskedLM = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForTokenClassification" }
 
 [families.esm3]
@@ -229,7 +279,7 @@ conversion_provenance = "Input: the pinned Biohub ESM3 checkpoint. Transformatio
 representative = "esm3_small"
 documentation = "docs/models.md#esm3"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esm3", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/esm3", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.esm3.modeling_esm3.FastESM3Config", AutoModel = "fastplms.models.esm3.modeling_esm3.FastESM3Model", AutoModelForSequenceClassification = "fastplms.models.esm3.modeling_esm3.FastESM3ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm3.modeling_esm3.FastESM3ForTokenClassification" }
 
 [families.e1]
@@ -256,7 +306,7 @@ conversion_provenance = "Input: the pinned Profluent-E1 checkpoint and tokenizer
 representative = "e1_150m"
 documentation = "docs/models.md#e1"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/e1", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/e1", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.e1.modeling_e1.E1Config", AutoModel = "fastplms.models.e1.modeling_e1.E1Model", AutoModelForMaskedLM = "fastplms.models.e1.modeling_e1.E1ForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.e1.modeling_e1.E1ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.e1.modeling_e1.E1ForTokenClassification" }
 
 [families.dplm]
@@ -282,7 +332,7 @@ conversion_provenance = "Input: the pinned official DPLM1 checkpoint. Transforma
 representative = "dplm_150m"
 documentation = "docs/models.md#dplm"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.dplm.modeling_dplm.DPLMConfig", AutoModel = "fastplms.models.dplm.modeling_dplm.DPLMModel", AutoModelForMaskedLM = "fastplms.models.dplm.modeling_dplm.DPLMForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.dplm.modeling_dplm.DPLMForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.dplm.modeling_dplm.DPLMForTokenClassification" }
 
 [families.dplm2]
@@ -307,7 +357,7 @@ conversion_provenance = "Input: the pinned official DPLM2 checkpoint. Transforma
 representative = "dplm2_150m"
 documentation = "docs/models.md#dplm2"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm2", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm2", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.dplm2.modeling_dplm2.DPLM2Config", AutoModel = "fastplms.models.dplm2.modeling_dplm2.DPLM2Model", AutoModelForMaskedLM = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForTokenClassification" }
 tokenizer_class = "fastplms.models.dplm2.tokenization_dplm2.DPLM2Tokenizer"
 
@@ -334,7 +384,7 @@ representative = "ankh_base"
 documentation = "docs/models.md#ankh"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
 requires_complete_weight_publication = false
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/ankh", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/ankh", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.ankh.modeling_ankh.FastAnkhConfig", AutoModel = "fastplms.models.ankh.modeling_ankh.FastAnkhModel", AutoModelForMaskedLM = "fastplms.models.ankh.modeling_ankh.FastAnkhForMaskedLMExtension", AutoModelForSeq2SeqLM = "fastplms.models.ankh.modeling_ankh.FastAnkhForConditionalGeneration", AutoModelForSequenceClassification = "fastplms.models.ankh.modeling_ankh.FastAnkhForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.ankh.modeling_ankh.FastAnkhForTokenClassification" }
 
 [families.boltz2]
@@ -358,7 +408,7 @@ conversion_provenance = "Input: the pinned official Boltz2 checkpoint. Transform
 representative = "boltz2"
 documentation = "docs/models.md#boltz2"
 test_tiers = ["structure", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "models/boltz"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "models/boltz"]
 auto_map = { AutoConfig = "fastplms.models.boltz.modeling_boltz2.Boltz2Config", AutoModel = "fastplms.models.boltz.modeling_boltz2.Boltz2Model" }
 
 [families.esmfold]
@@ -383,7 +433,7 @@ conversion_provenance = "Input: the pinned native Meta ESMFold checkpoint plus i
 representative = "esmfold"
 documentation = "docs/models.md#esmfold"
 test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/classification_probe.py", "models/esmfold"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/_esm_rotary.py", "models/classification_probe.py", "models/esmfold"]
 auto_map = { AutoConfig = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmFoldConfig", AutoModel = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForProteinFolding", AutoModelForSequenceClassification = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForTokenClassification" }
 
 [families.esmfold2]
@@ -410,7 +460,7 @@ conversion_provenance = "Input: each pinned Biohub ESMFold2 checkpoint and its s
 representative = "esmfold2"
 documentation = "docs/esmfold2.md"
 test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/classification_probe.py", "models/_esm_rotary.py", "models/esmfold2", "models/esm_plusplus", "models/ttt.py"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "atomic_files.py", "digests.py", "json_files.py", "models.toml", "models/__init__.py", "attention", "embeddings", "features","models/classification_probe.py", "models/_esm_rotary.py", "models/esmfold2", "models/esm_plusplus", "models/ttt.py"]
 auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2.ESMFold2Model", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ForTokenClassification" }
 
 [[models]]
@@ -418,7 +468,7 @@ id = "esm2_8m"
 family = "esm2"
 size_category = "small"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm2_8m.json=sha256:6975e86d1d8f27488bf2a676551feaa48cc19254c9d24b6acb09198122745609", tensors = "tests/goldens/esm2_8m.safetensors=sha256:b40217566c33c71988d28869de353be54a3b3ebfc21fdfd29056e88cf7e99f4c" }
+official_golden = { metadata = "tests/goldens/esm2_8m.json=sha256:06b3ccbf45a46a7aed3833b49503567f7534ffe60e6b3e49fcfcf17d04b7237e", tensors = "tests/goldens/esm2_8m.safetensors=sha256:d08a7572cbef20b8b19b545bcb0427b7e9ae19b986d015c559cd5a3a2cfc8aa4" }
 fast_repo = "Synthyra/ESM2-8M"
 fast_revision = "185ecbd45665d050a8dae326d91886d330c5f9d0"
 fast_files = [
@@ -457,7 +507,7 @@ id = "esm2_35m"
 family = "esm2"
 size_category = "small"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm2_35m.json=sha256:e919d3ce6d20b6a942d27d92323814ae7594a0129dc9c4de27c5053e96675bcd", tensors = "tests/goldens/esm2_35m.safetensors=sha256:c9b8bb616cf884fb7744521a2fcc6eed23586342d11241e6c9ef16454ec31e17" }
+official_golden = { metadata = "tests/goldens/esm2_35m.json=sha256:61cd24fd91ef2dc49cbc0f97b14c0d5c97849eca6f36ed8c294295147909cb38", tensors = "tests/goldens/esm2_35m.safetensors=sha256:4f82d10286e16041c2f23365dfdd8508b911864633287dd76580425527c2d922" }
 fast_repo = "Synthyra/ESM2-35M"
 fast_revision = "37ab9f56b41e365b3bd9e25d6fefe9150fd910f0"
 fast_files = [
@@ -496,7 +546,7 @@ id = "esm2_150m"
 family = "esm2"
 size_category = "medium"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm2_150m.json=sha256:c04c93486024ba0fa1c81fbfbe92ee79d1d4c7f1cfcc2c9886728522f752feab", tensors = "tests/goldens/esm2_150m.safetensors=sha256:c03fe9916dba137b452a6bbe944c7dc414db4019a6f0921e87b92d4bb6a8a42f" }
+official_golden = { metadata = "tests/goldens/esm2_150m.json=sha256:864e1f9c1d8e939e8da050b8617b8a7d6983fbb11a5911275d1571d08e460a79", tensors = "tests/goldens/esm2_150m.safetensors=sha256:20ad986b8e6e09f36d0158f5939d1914e0809b84329695028d767aab93338621" }
 fast_repo = "Synthyra/ESM2-150M"
 fast_revision = "979e0880dfc9e0c0080839b83d9d2dc05b92786a"
 fast_files = [
@@ -535,7 +585,7 @@ id = "esm2_650m"
 family = "esm2"
 size_category = "large"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm2_650m.json=sha256:f18332172fcb3abf5dd2485fd55f5b0d193ad3b93a44cc744e0d02817c927477", tensors = "tests/goldens/esm2_650m.safetensors=sha256:c3a66b75add03628e62e238cb63da6a9e4d321f8160e84bdf2a131c096977f86" }
+official_golden = { metadata = "tests/goldens/esm2_650m.json=sha256:b57f18c8803e68e6c621f479fd7a70596fbc416fc4a17d2ab91c639b480006c4", tensors = "tests/goldens/esm2_650m.safetensors=sha256:261a8b71c4b7c90b1f294031558ee0bf00d77a98352021eb78771200c8d40548" }
 fast_repo = "Synthyra/ESM2-650M"
 fast_revision = "ca0718a5d52b80d5c60dd76860e55e061a95fb0a"
 fast_files = [
@@ -574,7 +624,7 @@ id = "esm2_3b"
 family = "esm2"
 size_category = "xlarge"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm2_3b.json=sha256:5043b2333c57a34d54fac53916722d1acb4b6fd50395b9abafa805435b184a48", tensors = "tests/goldens/esm2_3b.safetensors=sha256:dfd5a8cb05d3e814a080185c4808c8e7ec2277f070f395562fcfbe4376789e4e" }
+official_golden = { metadata = "tests/goldens/esm2_3b.json=sha256:aa441786329889d43983811cac155cabd37b38cfd47bc90d415732a295c33c6d", tensors = "tests/goldens/esm2_3b.safetensors=sha256:61080e55e07db4a562a19e9b6b662e71a9fd3d8fee357aa3df00697d2fa58e33" }
 notes = "The pinned default SDPA BF16 path uses a checkpoint-specific numeric calibration: relative L2 target/hard limit 0.06/0.07, relative Q99.9 0.15/0.18, first-percentile residue cosine 0.994/0.992, and pooled cosine 0.998/0.997. Exact state identity and the global logits-distribution contract remain required."
 fast_repo = "Synthyra/ESM2-3B"
 fast_revision = "ff89d0180f414ab9c677219a25da79bf09185456"
@@ -617,7 +667,7 @@ id = "esmc_small"
 family = "esm_plusplus"
 size_category = "medium"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esmc_small.json=sha256:bb02652cf3cc484756b98ffa4ba55ed4c55870d2cea3342adb1d920ba9dfe10a", tensors = "tests/goldens/esmc_small.safetensors=sha256:03378d0f0fdd8161178ebb2c1f0da1b9776a726c8e8d3a10c009808a24de5654" }
+official_golden = { metadata = "tests/goldens/esmc_small.json=sha256:15cbbf909a12b5b81ff47b5f142babb75ebf9b3d85673905f2320c2a119ae347", tensors = "tests/goldens/esmc_small.safetensors=sha256:98219356e2845cbd10a80715b19c95632956d40567a7a993409c426b56040b06" }
 notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration."
 fast_repo = "Synthyra/ESMplusplus_small"
 fast_revision = "46c5f7d562e47d4c14165b424c71ab7db008e6fb"
@@ -643,7 +693,7 @@ id = "esmc_large"
 family = "esm_plusplus"
 size_category = "large"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esmc_large.json=sha256:7a4d614f67b6fde417f3fd89f61e7ec442ae284769734b2b73e14945a816a8fd", tensors = "tests/goldens/esmc_large.safetensors=sha256:e13302df4cf7e8381552f1043a8fd0f31f3e0d50b2ab6009fb86b7940ae8ff79" }
+official_golden = { metadata = "tests/goldens/esmc_large.json=sha256:86ffe179a9feafd067793f440d25e3606f184f0d32063aeaeb1842b14aa55d3b", tensors = "tests/goldens/esmc_large.safetensors=sha256:2b1ce3de5a7a27f171055f5c7aec4e7e50f0e809a1dc46508aaa0b472365015d" }
 notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration."
 fast_repo = "Synthyra/ESMplusplus_large"
 fast_revision = "f813401638b3fddab09748aec1ad2bf537aa4208"
@@ -669,7 +719,7 @@ id = "esmc_6b"
 family = "esm_plusplus"
 size_category = "xlarge"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esmc_6b.json=sha256:e229d938719782f280fab22dfc4c43e86109fdb0cc523631168c5a491afaace3", tensors = "tests/goldens/esmc_6b.safetensors=sha256:a948945e985c7deaca7be8b7eed09c0a9521a2af3f2b10fc2ec7a7d2a0f99ada" }
+official_golden = { metadata = "tests/goldens/esmc_6b.json=sha256:7e044e7e8d106d43169c4148521c54d0d09bb989624543d648247a110f1907d4", tensors = "tests/goldens/esmc_6b.safetensors=sha256:46139de6b28ab4f9e3244fb6ef3f259fc5e6518a3825f0fb4fd4a4fece6da8e6" }
 notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration."
 fast_repo = "Synthyra/ESMplusplus_6B"
 fast_revision = "0d579cce3b0f09efa6b3baddf6cc3fd8c9b616c8"
@@ -681,6 +731,7 @@ fast_files = [
   "model-00004-of-00006.safetensors=sha256:e46c6113c89c6f3e9b072c1bef02d763a625c37bcd8f9da2ed9363891c9a0758",
   "model-00005-of-00006.safetensors=sha256:6d92cb2bf9791de644de2ae86f8523d802ac3b4aaabfff0716ab6c2b97f6fb14",
   "model-00006-of-00006.safetensors=sha256:5fc1a8632490bb34162823c35d0d591337b9e4195b22cc0560741397a6e9d0b3",
+  "model.safetensors.index.json=git-sha1:f30f0b6b35e11d09a516c078b0847ee023bd1817",
   "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
   "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71",
   "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756",
@@ -706,7 +757,7 @@ family = "esm3"
 tokenizer_source = "esmc_small"
 size_category = "large"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/esm3_small.json=sha256:5470e8596cbba0e2882647eccbc53c36d8b48b0f3947d1fe0bcea68da1078c32", tensors = "tests/goldens/esm3_small.safetensors=sha256:d957922f810c9ab4c557d80d5aaaf6a3aab79a5a45e4638012a634a4134803b1" }
+official_golden = { metadata = "tests/goldens/esm3_small.json=sha256:e476bad7ccdfa2f908261a03e20116c0854eece292a553196f940b594088204d", tensors = "tests/goldens/esm3_small.safetensors=sha256:251e050926a1f2426401bac6b93b4cb00041d0b0a5c8a458d0ef0e6a5d0d87c3" }
 fast_repo = "Synthyra/ESM3_small"
 fast_revision = "7ddb5a740f9e5f93933eb6410c0ee8684bc63ec1"
 fast_files = [
@@ -732,7 +783,7 @@ id = "e1_150m"
 family = "e1"
 size_category = "small"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/e1_150m.json=sha256:701a64a6ab1a2fec5a427555b6af96232526c15cb3d5b4dc7fb253ac8f20b922", tensors = "tests/goldens/e1_150m.safetensors=sha256:6558bc8f1a7b20629eaaaa6f72601d0c2cdb859a5dc13595549b1773b6e2de41" }
+official_golden = { metadata = "tests/goldens/e1_150m.json=sha256:e7520331135fd4c98f52f21c470dbba5cb5f4defb418494c63199fb000621e65", tensors = "tests/goldens/e1_150m.safetensors=sha256:03a2e93e7b3e54b12f99eea7678ff80828922f92c3683d4909346a05d74815bc" }
 fast_repo = "Synthyra/Profluent-E1-150M"
 fast_revision = "7c5f3bbf697226a2e0900db7a100f9201774a907"
 fast_files = [
@@ -751,7 +802,7 @@ id = "e1_300m"
 family = "e1"
 size_category = "medium"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/e1_300m.json=sha256:d3478f3f5957a0e0377864074dde0107de890019f96cb63548ee17ffb8f3ec3a", tensors = "tests/goldens/e1_300m.safetensors=sha256:92778b9ef95a803ddc84b3e3ca764c59e045872a94bcff0eb0cd47647732c188" }
+official_golden = { metadata = "tests/goldens/e1_300m.json=sha256:d97637bf789fed61b4ad575f904f37befcea42e9436e8d841e8f8f2cff1d1b17", tensors = "tests/goldens/e1_300m.safetensors=sha256:d125ea866899d2788330d6877d77c48c69febbf7f01d629f4107f9647cfa13ff" }
 fast_repo = "Synthyra/Profluent-E1-300M"
 fast_revision = "5ef52c0ad2ae2578f40622696b763523810e8e26"
 fast_files = [
@@ -770,7 +821,7 @@ id = "e1_600m"
 family = "e1"
 size_category = "large"
 generation_contract = "not_applicable"
-official_golden = { metadata = "tests/goldens/e1_600m.json=sha256:914be191c28141c1f84535cdb69ead0588a2057bb19d46c5bc7f3891a3d6739e", tensors = "tests/goldens/e1_600m.safetensors=sha256:22ed8417a4651ded255099f6d15c63c2c40552e700d2b0470d1adfde3a39c513" }
+official_golden = { metadata = "tests/goldens/e1_600m.json=sha256:0c9afc83b96ed8ad7f339a26df4ba5853e7994ad54dd3cd837deb9869c0cd2a9", tensors = "tests/goldens/e1_600m.safetensors=sha256:d470d6833f0b38bec070aa37e9f3c49717ea82d31c794af843550bb7e56253a6" }
 fast_repo = "Synthyra/Profluent-E1-600M"
 fast_revision = "6c8bf0ec83b0e0178677c528b101efffd0677742"
 fast_files = [
@@ -789,7 +840,7 @@ id = "dplm_150m"
 family = "dplm"
 size_category = "small"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/dplm_150m.json=sha256:3228551fe3bed951db9ec97347143ec4462ce7c221ac240b7ce7730948c1dc1f", tensors = "tests/goldens/dplm_150m.safetensors=sha256:392992235195beed97ab8359b90a2e11e52f4326606f99a471447bed81d146bd" }
+official_golden = { metadata = "tests/goldens/dplm_150m.json=sha256:541167ffcd12b2d7e101046b2d17e800c70ba5db1b561212575301f2d1c407f8", tensors = "tests/goldens/dplm_150m.safetensors=sha256:f8e4ddd580d4708dad2d7da002735b2a01d69ece478c46c4932eb2f3933dc836" }
 fast_repo = "Synthyra/DPLM-150M"
 fast_revision = "90ba742754151a774f3b7ed580170d0a76b3e69d"
 fast_files = [
@@ -814,7 +865,7 @@ id = "dplm_650m"
 family = "dplm"
 size_category = "large"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/dplm_650m.json=sha256:bf58d0ce73aaac7e6fb1923ef3d9adad67122df2a3dd414c3229488ef9587a6d", tensors = "tests/goldens/dplm_650m.safetensors=sha256:073f0a6abea7e48f28c2d921ff8329a28e22627f01979277cb324908a01b3378" }
+official_golden = { metadata = "tests/goldens/dplm_650m.json=sha256:2bd482c94b1b6b3ac901b5ff9b06e2d189cda1a8fe2f8ff32e5d52b64c2c446b", tensors = "tests/goldens/dplm_650m.safetensors=sha256:70e2c46ed722942d3b628b556b92c9994688cb4853b278eb1e40a3dfbbd9dfc4" }
 fast_repo = "Synthyra/DPLM-650M"
 fast_revision = "05dc16d97c5c028aed924c9ed681cee4ab609760"
 fast_files = [
@@ -839,7 +890,7 @@ id = "dplm_3b"
 family = "dplm"
 size_category = "xlarge"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/dplm_3b.json=sha256:a5b6df8b9c7b371976892ec1d6c45581a32ad3a6325c6c0a0b3267012848c8ed", tensors = "tests/goldens/dplm_3b.safetensors=sha256:75b0a0854fc391133920b0feaaeb8f69ab7568a88b3759627aca1556c4338c1e" }
+official_golden = { metadata = "tests/goldens/dplm_3b.json=sha256:0f8d8fb7df562c30abf730abe55db4a8c13d40877dd21772375fe1f34030159a", tensors = "tests/goldens/dplm_3b.safetensors=sha256:6942218079232a8185b1ae6e1578474c7784722e28f528b76ec89ca77777e2b9" }
 fast_repo = "Synthyra/DPLM-3B"
 fast_revision = "7d764dd3d70ecf1ac0e64693de64a0064aacac65"
 fast_files = [
@@ -869,7 +920,7 @@ id = "dplm2_150m"
 family = "dplm2"
 size_category = "small"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/dplm2_150m.json=sha256:d269de779ea1503de72c77e7b2e6224afc9797bd945b40c571ff6faec782e4aa", tensors = "tests/goldens/dplm2_150m.safetensors=sha256:17fc26600938ba5364b8ecb96750786d33e9f92bcd4ea4df3e12a389340748eb" }
+official_golden = { metadata = "tests/goldens/dplm2_150m.json=sha256:5f2c496e4b557faf70bdde03b531f753fc7f20c5868bab79dcf378ee358781a1", tensors = "tests/goldens/dplm2_150m.safetensors=sha256:7390559bc54f27f08972c8f0adcf17954065450d937b3984303c1ae1ec003855" }
 artifact_source = "official"
 canonical_state_sha256 = "82e1751f59052b8de72b082517557db47947e8d9b4ac2f11278369e6c0cbf001"
 fast_repo = "Synthyra/DPLM2-150M"
@@ -896,7 +947,7 @@ id = "dplm2_650m"
 family = "dplm2"
 size_category = "large"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/dplm2_650m.json=sha256:d9a7548f9af657a72d441ca70f27379863724fcce8ddd3da4f672104b7bfb772", tensors = "tests/goldens/dplm2_650m.safetensors=sha256:c4e0e467c252c3ac813363d2d4b17a5e3bd99e75fad315e76d97689b4655ddac" }
+official_golden = { metadata = "tests/goldens/dplm2_650m.json=sha256:af7b8c4fab48773e1ddd2cfd941cf3d518ea27ed8519f467ddf0da788a0aa6be", tensors = "tests/goldens/dplm2_650m.safetensors=sha256:c45f4ca267183b11dd06244ed3c3f39cf5b6139f778733fa63a6e2b5a46b0678" }
 artifact_source = "official"
 canonical_state_sha256 = "cba76b6602d2258de9fffff953b608d93cb8ef4a9e89b0bbd27e160c81e78bb4"
 fast_repo = "Synthyra/DPLM2-650M"
@@ -925,7 +976,7 @@ size_category = "xlarge"
 # The pinned public sampler fails before generation because cls_token_id is None.
 # State, tokenizer, and inference parity remain required for this checkpoint.
 generation_contract = "official_unavailable"
-official_golden = { metadata = "tests/goldens/dplm2_3b.json=sha256:d6e0e02af53b13cb129192f06e264758aa21c9ebf4ee82411cf67037082d2329", tensors = "tests/goldens/dplm2_3b.safetensors=sha256:838b11824d08f83bcb0c0b3268e579f3a87dbfb965370cfe5c3f8793b96b1964" }
+official_golden = { metadata = "tests/goldens/dplm2_3b.json=sha256:ded7978cf84da1d56d8b4693418030ec3b0086bf6466be2242fb0ce424296b9f", tensors = "tests/goldens/dplm2_3b.safetensors=sha256:7490f771f8b8f8137b333f98d6b7dbd2cb48708c8649d4459279a6711b9e5870" }
 notes = "The pinned official DPLM2-3B sampler fails before generation, so live generation equivalence cannot be established for this checkpoint. State, tokenizer, and inference parity remain required."
 artifact_source = "official"
 canonical_state_sha256 = "8c46ec09115dbe6cbfb91d94ab5e906369d57e27fe620a7741c6f8cb1b6ca890"
@@ -958,7 +1009,7 @@ id = "ankh_base"
 family = "ankh"
 size_category = "medium"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/ankh_base.json=sha256:ebce8d7de821827ee995789c9b38d79252d3b2f76888130b0a8a7eedafaefe2b", tensors = "tests/goldens/ankh_base.safetensors=sha256:f0e78aa15d11749e0c64ff57f9e88c51cec6538a0adf8951f839df70cc708b65" }
+official_golden = { metadata = "tests/goldens/ankh_base.json=sha256:21f129b7b71c026cbd3b3fdb712af06d5b6f256b487ef8b01fcc506d2bc6c7b6", tensors = "tests/goldens/ankh_base.safetensors=sha256:449db3638117d5b6027380fdf4b05698bb156024b2b3f7e0ceb295a86e5ec756" }
 notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
 artifact_source = "official"
 canonical_state_sha256 = "cdd8d30d88e5bf41f44e1eef4470d8e46607aba5f7c7c805b06c035b89c8c16f"
@@ -987,7 +1038,7 @@ id = "ankh_large"
 family = "ankh"
 size_category = "large"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/ankh_large.json=sha256:59492518b021de5cfaea87d672c9448c8558e99a3443ba2cc7ab544963196ecb", tensors = "tests/goldens/ankh_large.safetensors=sha256:3fb8d3ac27716d15a9ea92aeef6acf2b977bcc887d9b535000539e523673459b" }
+official_golden = { metadata = "tests/goldens/ankh_large.json=sha256:8f269215ee29887661abb6c139139dd1782f5112b94eb63be968282a7f050e1b", tensors = "tests/goldens/ankh_large.safetensors=sha256:4a7aea66e91cf880b0b839495d7704831725b672fe169f8e7258b5643d77aff9" }
 notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
 artifact_source = "official"
 canonical_state_sha256 = "e498a2e9aea76ef784cbe3e596c6b3f5e9a40e209ad837f7e3207099e4d74483"
@@ -1017,7 +1068,7 @@ id = "ankh2_large"
 family = "ankh"
 size_category = "large"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/ankh2_large.json=sha256:e8df38994ca1a1e0c598ace34a0b257b264937e4fdbb01bc41544985116b02a4", tensors = "tests/goldens/ankh2_large.safetensors=sha256:25fe1569f55c635fab8fa49c1d62a889a35a2a738bad921f5764a85b58fd4b5d" }
+official_golden = { metadata = "tests/goldens/ankh2_large.json=sha256:bdafcd179bc055ea223de43228801b1b6c35196741ac4cc4cca0571e3291e6dc", tensors = "tests/goldens/ankh2_large.safetensors=sha256:44d5fc5c74a0d166ceac54d12505a109f8338cb7090af2f37e39bd78545a3c52" }
 notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
 artifact_source = "official"
 canonical_state_sha256 = "597c4fe2fa8711f11a25317905f1d62fa92905e55fdd5c0a79614cd9c9d2bca3"
@@ -1049,7 +1100,7 @@ id = "ankh3_large"
 family = "ankh"
 size_category = "large"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/ankh3_large.json=sha256:2e5bb05b3baa5baa78f61fef7d2a2c669b0da5dbfaf6b50b12abd3e17253a961", tensors = "tests/goldens/ankh3_large.safetensors=sha256:e5c494ac418e0a2fe7bdad1376676d48960d58ec9e044d19bfffccb8c3288513" }
+official_golden = { metadata = "tests/goldens/ankh3_large.json=sha256:ab3a671260bdad481635d4a4be1b8072d5a2e02e3df1178202713c9a7e56af49", tensors = "tests/goldens/ankh3_large.safetensors=sha256:cc55274281e87d22cf7828117678df2e37a1ece543f5b401b65a0cd6f8b6fa27" }
 notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
 artifact_source = "official"
 canonical_state_sha256 = "60acb7ef86e85dc0c51fc1edf4c8e69a0480049723b6b2c95e6e9faa720c112a"
@@ -1083,7 +1134,7 @@ id = "ankh3_xl"
 family = "ankh"
 size_category = "xlarge"
 generation_contract = "required"
-official_golden = { metadata = "tests/goldens/ankh3_xl.json=sha256:66bb12e033e4163be225d636108a479393228a4f5061015c8af114e766c3c486", tensors = "tests/goldens/ankh3_xl.safetensors=sha256:72d34567d0228cb6f1ee701c578ed4039fead4346e3f161a52e0e74df28dc8ae" }
+official_golden = { metadata = "tests/goldens/ankh3_xl.json=sha256:a2d302227ee616698af18502d03e2b6b136d589134058d82b018520f18498431", tensors = "tests/goldens/ankh3_xl.safetensors=sha256:d18c086a4849b220761019f219590fe0c1c6f17698c365ccb0a5a8804446286c" }
 notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head. The official PyTorch shard index is deliberately excluded: the builder verifies every declared source shard directly and writes a new canonical safetensors index."
 artifact_source = "official"
 canonical_state_sha256 = "dd2188e0d2ca65232135714eef6de394239734d843ddae4928c7398685d858e7"
@@ -1176,7 +1227,7 @@ family = "esmfold2"
 size_category = "structure"
 generation_contract = "not_applicable"
 msa_conditioning = true
-official_golden = { metadata = "tests/goldens/esmfold2.json=sha256:f6e0ed1ec400b9a0fcc817db51774be968dc454b7a32645a07c479e42423ab20", tensors = "tests/goldens/esmfold2.safetensors=sha256:e4d6be4344c528e26b13f79a9303549e3de7e582da195c0078db3ce957fad420" }
+official_golden = { metadata = "tests/goldens/esmfold2.json=sha256:60b9e8827615a73acd89ac23f13e3d0a412c17879fbf3308444546e6bbb4ff03", tensors = "tests/goldens/esmfold2.safetensors=sha256:a1a9f3b7a9f7e36ef6a7077d76fa1b1c8b5529ed42ac2cfa485797162ce866fb" }
 fast_repo = "Synthyra/ESMFold2"
 fast_revision = "cd5a0927cec585a778d983b99a8db23d2e9b281e"
 fast_files = [
@@ -1196,7 +1247,7 @@ family = "esmfold2"
 size_category = "structure"
 generation_contract = "not_applicable"
 msa_conditioning = false
-official_golden = { metadata = "tests/goldens/esmfold2_fast.json=sha256:091b004c0b330217b59c12acd6da3d6edaf91e48d95f6d5f40fc20399cef9478", tensors = "tests/goldens/esmfold2_fast.safetensors=sha256:6e2e1cd07401538b4d9df994f82abe7a5b38a01e8d1ee26681e1216d44a81990" }
+official_golden = { metadata = "tests/goldens/esmfold2_fast.json=sha256:659bb338be4584787faff619d1cbcb8766aece7b5facf8a84470e850299b662f", tensors = "tests/goldens/esmfold2_fast.safetensors=sha256:8da91024a2cc63984585267a18f595892ebe489b2a05c4f79a994c3ac5de2f40" }
 fast_repo = "Synthyra/ESMFold2-Fast"
 fast_revision = "407875bfcaa42552bfcb25acd67ee1888b790170"
 fast_files = [
@@ -1216,7 +1267,7 @@ family = "esmfold2"
 size_category = "structure"
 generation_contract = "not_applicable"
 msa_conditioning = true
-official_golden = { metadata = "tests/goldens/esmfold2_experimental_cutoff2025.json=sha256:cfd0e35b2bc468a0dc4f614d3acfa2fce004f96e9ae2433256ed095b829d55cc", tensors = "tests/goldens/esmfold2_experimental_cutoff2025.safetensors=sha256:9347466bbe803b6f5dc82e3356ca6cbbf2c2edd8765f9fd273385bda255019f6" }
+official_golden = { metadata = "tests/goldens/esmfold2_experimental_cutoff2025.json=sha256:2e42d958f7a99ca8edaf0d9fac8e1a58b78200df652116b7bc102ac799123633", tensors = "tests/goldens/esmfold2_experimental_cutoff2025.safetensors=sha256:8168f8692c4932e1160e733b15210836120223615f6190fb0ecfa2f940c2420f" }
 fast_repo = "Synthyra/ESMFold2-Experimental-Cutoff2025"
 fast_revision = "632ff4a9e68f1de78ee956a613267bdcdb5b354d"
 fast_files = [
@@ -1237,7 +1288,7 @@ family = "esmfold2"
 size_category = "structure"
 generation_contract = "not_applicable"
 msa_conditioning = false
-official_golden = { metadata = "tests/goldens/esmfold2_experimental_fast_cutoff2025.json=sha256:1d0b2da4f1579243f37ae04bd4b834b747005cd8e8e7665e00d088123c43afd9", tensors = "tests/goldens/esmfold2_experimental_fast_cutoff2025.safetensors=sha256:516e216d05d7e6bee59e77126d3e595e2bb7821929433f00c259c5d5241964bb" }
+official_golden = { metadata = "tests/goldens/esmfold2_experimental_fast_cutoff2025.json=sha256:d41c54ca28270fc1c1b95485c5009990bd2c215d35ec5cc2c2fa611d096bad21", tensors = "tests/goldens/esmfold2_experimental_fast_cutoff2025.safetensors=sha256:97b58525aefd21c23ec6cb2defa3b3b41982c872dfe70d94306648edbe40a5b9" }
 fast_repo = "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025"
 fast_revision = "8f022c2514a6c32692aaca078a8391d6bc6c4bac"
 fast_files = [
@@ -1254,6 +1305,7 @@ auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFo
 
 [[models]]
 id = "esmfold2_300"
+official_golden = { metadata = "tests/goldens/esmfold2_300.json=sha256:d4fd12f27352c53582bb40a8d85eab76d6d98f647ce7660a6f891a1dfe68c039", tensors = "tests/goldens/esmfold2_300.safetensors=sha256:40a954ff14edc7c2bff95a252241f08db4cae3821ca9b92ae216775cceed8e8b" }
 confidence_adaptation = { release = "v1", head_sha256 = "40fd7f3d82fcefe8ad20ab2b32a37a68a84b54a527a4bce6eb9437bad2b77e31", base_weight_sha256 = "44d6797c5efebf24753d502b40950e0874871c96ceea14f2d7f7e39cebac67fd", donor_repo = "biohub/ESMFold2-Experimental-Fast-Cutoff2025", donor_revision = "74b88548bf19688b8727432db0d698cb2e1d8783", donor_weight_sha256 = "4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f", training_url = "https://wandb.ai/lhallee/fastplms-confidence/runs/9558b6d23daf", evaluation_url = "https://huggingface.co/datasets/Synthyra/FastPLMs-artifacts/tree/309f353b07e0e46de4d77a5266d4eddd695538e3/confidence-v2/v2-reproduction-20260922/public/evaluation/esmfold2_300", evidence_path = "docs/evidence/confidence/esmfold2_300-v1.json", frozen_base = { repo = "Synthyra/ESMFold2-300", revision = "a38a62ae930d157484b331c2bf4241684573adba", files = ["config.json=git-sha1:47ec20cf8b234c3b41d6f3ae1bdfe95d4eb4849e", "model.safetensors=sha256:44d6797c5efebf24753d502b40950e0874871c96ceea14f2d7f7e39cebac67fd"] } }
 family = "esmfold2"
 size_category = "structure"
@@ -1273,6 +1325,7 @@ backbone = { repo = "biohub/ESMC-300M-1500000", revision = "56803b6378b82e16c3b2
 
 [[models]]
 id = "esmfold2_600"
+official_golden = { metadata = "tests/goldens/esmfold2_600.json=sha256:ea249ff118975f29979143ea081cfa1c0fb60cb1fa393dd99a12caf413524230", tensors = "tests/goldens/esmfold2_600.safetensors=sha256:69bbb8d53a816b29e979e4469728df0a6b662e4d0be81a76068b36974d2b909e" }
 confidence_adaptation = { release = "v1", head_sha256 = "e84726a050722e1b722712c87d17a5388bd3699e2520e4a59abb1d828dfb8de7", base_weight_sha256 = "11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602", donor_repo = "biohub/ESMFold2-Experimental-Fast-Cutoff2025", donor_revision = "74b88548bf19688b8727432db0d698cb2e1d8783", donor_weight_sha256 = "4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f", training_url = "https://wandb.ai/lhallee/fastplms-confidence/runs/820d2cfa56c0", evaluation_url = "https://huggingface.co/datasets/Synthyra/FastPLMs-artifacts/tree/6e62186cd36b9047cc4691980076be9f76482192/confidence-v2/v2-reproduction-20260922/public/evaluation/esmfold2_600", evidence_path = "docs/evidence/confidence/esmfold2_600-v1.json", frozen_base = { repo = "Synthyra/ESMFold2-600", revision = "71c67d0b2b73dc245ea7c3cc0d0476439a882d08", files = ["config.json=git-sha1:8e271837cbdada96c4974c8e543f84065e0f06f1", "model.safetensors=sha256:11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602"] } }
 family = "esmfold2"
 size_category = "structure"
@@ -1289,3 +1342,277 @@ notes = "Experimental Fast model with a frozen 600M ESM++ backbone, 24 folding b
 auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" }
 backbone_model = "esmc_large"
 backbone = { repo = "biohub/ESMC-600M-1500000", revision = "21af9cc429af76ebda6c48074fb624db4735aaaf", files = ["config.json=git-sha1:ec29f6009b21d710f64bf1c058f3a9710833d692", "model.safetensors=sha256:d6869f5ae0f11e5dc829b195e062e87cfcc2f851a08a5edbaf5d1083ae7f76cc", "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", "tokenizer_config.json=git-sha1:f49f57b24a1c93bd544974811e8ecbd61b7fae89", "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b"] }
+
+# Parity goldens, uploaded 2026-09-22 so the reference tensors live somewhere durable
+# rather than only inside the GitHub repository. Each sha256 is the one already recorded
+# in that model's official_golden.tensors, so the existing integrity check is unchanged.
+
+[[golden_artifacts]]
+id = "esm2_8m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm2_8m.safetensors"
+sha256 = "d08a7572cbef20b8b19b545bcb0427b7e9ae19b986d015c559cd5a3a2cfc8aa4"
+size = 537589
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esm2_35m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm2_35m.safetensors"
+sha256 = "4f82d10286e16041c2f23365dfdd8508b911864633287dd76580425527c2d922"
+size = 779509
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esm2_150m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm2_150m.safetensors"
+sha256 = "20ad986b8e6e09f36d0158f5939d1914e0809b84329695028d767aab93338621"
+size = 1021437
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esm2_650m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm2_650m.safetensors"
+sha256 = "261a8b71c4b7c90b1f294031558ee0bf00d77a98352021eb78771200c8d40548"
+size = 1989117
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esm2_3b"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm2_3b.safetensors"
+sha256 = "61080e55e07db4a562a19e9b6b662e71a9fd3d8fee357aa3df00697d2fa58e33"
+size = 3924485
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmc_small"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmc_small.safetensors"
+sha256 = "98219356e2845cbd10a80715b19c95632956d40567a7a993409c426b56040b06"
+size = 1165354
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmc_large"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmc_large.safetensors"
+sha256 = "2b1ce3de5a7a27f171055f5c7aec4e7e50f0e809a1dc46508aaa0b472365015d"
+size = 1383082
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmc_6b"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmc_6b.safetensors"
+sha256 = "46139de6b28ab4f9e3244fb6ef3f259fc5e6518a3825f0fb4fd4a4fece6da8e6"
+size = 2979762
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esm3_small"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esm3_small.safetensors"
+sha256 = "251e050926a1f2426401bac6b93b4cb00041d0b0a5c8a458d0ef0e6a5d0d87c3"
+size = 2398877
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "e1_150m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/e1_150m.safetensors"
+sha256 = "03a2e93e7b3e54b12f99eea7678ff80828922f92c3683d4909346a05d74815bc"
+size = 958851
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "e1_300m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/e1_300m.safetensors"
+sha256 = "d125ea866899d2788330d6877d77c48c69febbf7f01d629f4107f9647cfa13ff"
+size = 1258379
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "e1_600m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/e1_600m.safetensors"
+sha256 = "d470d6833f0b38bec070aa37e9f3c49717ea82d31c794af843550bb7e56253a6"
+size = 1557899
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm_150m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm_150m.safetensors"
+sha256 = "f8e4ddd580d4708dad2d7da002735b2a01d69ece478c46c4932eb2f3933dc836"
+size = 1021437
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm_650m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm_650m.safetensors"
+sha256 = "70e2c46ed722942d3b628b556b92c9994688cb4853b278eb1e40a3dfbbd9dfc4"
+size = 1989117
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm_3b"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm_3b.safetensors"
+sha256 = "6942218079232a8185b1ae6e1578474c7784722e28f528b76ec89ca77777e2b9"
+size = 3924485
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm2_150m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm2_150m.safetensors"
+sha256 = "7390559bc54f27f08972c8f0adcf17954065450d937b3984303c1ae1ec003855"
+size = 26826946
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm2_650m"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm2_650m.safetensors"
+sha256 = "c45f4ca267183b11dd06244ed3c3f39cf5b6139f778733fa63a6e2b5a46b0678"
+size = 28762314
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "dplm2_3b"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/dplm2_3b.safetensors"
+sha256 = "7490f771f8b8f8137b333f98d6b7dbd2cb48708c8649d4459279a6711b9e5870"
+size = 32633034
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "ankh_base"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/ankh_base.safetensors"
+sha256 = "449db3638117d5b6027380fdf4b05698bb156024b2b3f7e0ceb295a86e5ec756"
+size = 860722
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "ankh_large"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/ankh_large.safetensors"
+sha256 = "4a7aea66e91cf880b0b839495d7704831725b672fe169f8e7258b5643d77aff9"
+size = 1717818
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "ankh2_large"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/ankh2_large.safetensors"
+sha256 = "44d5fc5c74a0d166ceac54d12505a109f8338cb7090af2f37e39bd78545a3c52"
+size = 1717818
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "ankh3_large"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/ankh3_large.safetensors"
+sha256 = "cc55274281e87d22cf7828117678df2e37a1ece543f5b401b65a0cd6f8b6fa27"
+size = 1717818
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "ankh3_xl"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/ankh3_xl.safetensors"
+sha256 = "d18c086a4849b220761019f219590fe0c1c6f17698c365ccb0a5a8804446286c"
+size = 2860602
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold.safetensors"
+sha256 = "873b1b325a43d8e0f35f355c8914a2a9fe611cc48763875e9e6a22e09ec9ebcb"
+size = 179144
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2.safetensors"
+sha256 = "a1a9f3b7a9f7e36ef6a7077d76fa1b1c8b5529ed42ac2cfa485797162ce866fb"
+size = 2573976
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2_fast"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2_fast.safetensors"
+sha256 = "8da91024a2cc63984585267a18f595892ebe489b2a05c4f79a994c3ac5de2f40"
+size = 2573976
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2_experimental_cutoff2025"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2_experimental_cutoff2025.safetensors"
+sha256 = "8168f8692c4932e1160e733b15210836120223615f6190fb0ecfa2f940c2420f"
+size = 1771064
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2_experimental_fast_cutoff2025"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2_experimental_fast_cutoff2025.safetensors"
+sha256 = "97b58525aefd21c23ec6cb2defa3b3b41982c872dfe70d94306648edbe40a5b9"
+size = 1771064
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2_300"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2_300.safetensors"
+sha256 = "40a954ff14edc7c2bff95a252241f08db4cae3821ca9b92ae216775cceed8e8b"
+size = 865376
+offline_behavior = "requires_cached_verified_file"
+
+[[golden_artifacts]]
+id = "esmfold2_600"
+repository = "Synthyra/fastplms_parity_goldens"
+revision = "50b93a0fc2be21a36c2522118754c74c4a9631ad"
+path = "goldens/esmfold2_600.safetensors"
+sha256 = "69bbb8d53a816b29e979e4469728df0a6b662e4d0be81a76068b36974d2b909e"
+size = 865376
+offline_behavior = "requires_cached_verified_file"

fastplms/models/esm_plusplus/modeling_esm_plusplus.py

@@ -10,7 +10,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from contextlib import contextmanager
 from dataclasses import asdict, dataclass
 from functools import partial
@@ -416,7 +416,7 @@ class RotaryEmbedding(torch.nn.Module):
             and self._seq_len_cached >= token_count
             and cached.device == device
             and cached.dtype == dtype
-            and not (self.training and cached.is_inference())
+            and not (torch.is_grad_enabled() and cached.is_inference())  # autograd cannot save inference-created tables, in eval mode too
         )
 
     def _rotary_angles(
@@ -808,6 +808,7 @@ class TransformerOutput(ModelOutput):
     s_max: tuple[list[torch.Tensor], ...] | None = None
     sae_outputs: dict[str, torch.Tensor] | None = None
     sae_hidden_states: dict[int, torch.Tensor] | None = None
+    captured_hidden_states: dict[int, torch.Tensor] | None = None
 
 
 @dataclass
@@ -884,13 +885,53 @@ class TransformerStack(nn.Module):
         output_s_max: bool | None = False,
         esmfold2_hidden_states: bool = False,
         sae_layers: tuple[int, ...] = (),
+        capture_layers: tuple[int, ...] = (),
+        stop_after_layer: int | None = None,
+        stream_layers: tuple[int, ...] = (),
+        state_consumer: Callable[[int, torch.Tensor], None] | None = None,
+        assume_valid_mask: bool = False,
     ) -> TransformerOutput:
+        """Run the blocks and record the requested hidden states.
+
+        Hidden state ``i`` is the input to block ``i``, and state ``n_layers`` is the final
+        normalized state. ``capture_layers`` and ``sae_layers`` record states by that index.
+        ``stop_after_layer`` ends the pass once its state exists: no later block runs, and the
+        final norm runs only when that state is ``n_layers``, the default.
+        ``assume_valid_mask`` skips the check that every row keeps a valid key. That check reads a
+        device value and so stalls the host until the device is idle; a caller that builds its mask
+        from known lengths, as the canonical token executor does, sets it.
+        """
         # x: (b, l, d); attention_mask, sequence_id: (b, l)
+        n_layers = len(self.blocks)
+        depth = n_layers if stop_after_layer is None else stop_after_layer
+        if not 0 <= depth <= n_layers:
+            raise ValueError(
+                f"stop_after_layer must be a hidden-state index in 0..{n_layers}; "
+                f"received {stop_after_layer!r}."
+            )
+        capture_set = set(capture_layers)
+        stream_set = set(stream_layers)
+        if bool(stream_set) != (state_consumer is not None):
+            raise ValueError("Streaming layers and their state consumer must be supplied together.")
+        sae_layer_set = set(sae_layers)
+        # A full pass keeps the SAE contract unchanged; an early stop must reach every SAE state.
+        requested = capture_set | sae_layer_set if depth < n_layers else capture_set
+        requested = requested | stream_set
+        uncomputed = sorted(index for index in requested if not 0 <= index <= depth)
+        if uncomputed:
+            raise ValueError(
+                f"Hidden states {uncomputed} are outside the states 0..{depth} this call computes."
+            )
+        if depth < n_layers and (output_hidden_states or output_attentions or output_s_max):
+            raise ValueError(
+                "A stack that stops before its final state returns only captured hidden states; "
+                "output_hidden_states, output_attentions, and output_s_max need every block."
+            )
         hidden_states = () if output_hidden_states else None
         attentions = () if output_attentions else None
         full_s_max = () if output_s_max else None
-        sae_layer_set = set(sae_layers)
-        sae_hidden_states = {} if sae_layer_set else None
+        sae_hidden_states: dict[int, torch.Tensor] | None = {} if sae_layer_set else None
+        captured_hidden_states: dict[int, torch.Tensor] | None = {} if capture_set else None
         # 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.
@@ -902,6 +943,7 @@ class TransformerStack(nn.Module):
             device=x.device,
             dtype=x.dtype,
             output_attentions=bool(output_attentions),
+            validate_mask=not assume_valid_mask,
         )
         # A call that returns attention weights runs eager attention and needs no layout.
         flash_padding_layout = (
@@ -911,6 +953,8 @@ class TransformerStack(nn.Module):
         )
 
         for layer_index, block in enumerate(self.blocks):
+            if layer_index == depth:
+                break
             if output_hidden_states:
                 if hidden_states is None:
                     raise RuntimeError(
@@ -922,6 +966,10 @@ class TransformerStack(nn.Module):
                 hidden_states += (x,)
             if sae_hidden_states is not None and layer_index in sae_layer_set:
                 sae_hidden_states[layer_index] = x
+            if captured_hidden_states is not None and layer_index in capture_set:
+                captured_hidden_states[layer_index] = x
+            if layer_index in stream_set:
+                state_consumer(layer_index, x)
             if self.gradient_checkpointing and self.training:
                 x, attn_weights, s_max = self._gradient_checkpointing_func(
                     block.__call__,
@@ -949,12 +997,20 @@ class TransformerStack(nn.Module):
             if full_s_max is not None:
                 full_s_max += (s_max,)
 
-        last_hidden_state = self.norm(x)
-        if output_hidden_states:
-            hidden_states += (last_hidden_state,)
-        final_layer_index = len(self.blocks)
-        if sae_hidden_states is not None and final_layer_index in sae_layer_set:
-            sae_hidden_states[final_layer_index] = last_hidden_state
+        last_hidden_state: torch.Tensor | None = None
+        if depth == n_layers:
+            last_hidden_state = self.norm(x)  # (b, l, d)
+            if output_hidden_states:
+                hidden_states += (last_hidden_state,)
+        # State `depth` is the final normalized state, or, after an early stop, the input to
+        # block `depth`, which the final norm never touches.
+        deepest_state = x if last_hidden_state is None else last_hidden_state  # (b, l, d)
+        if sae_hidden_states is not None and depth in sae_layer_set:
+            sae_hidden_states[depth] = deepest_state
+        if captured_hidden_states is not None and depth in capture_set:
+            captured_hidden_states[depth] = deepest_state
+        if depth in stream_set:
+            state_consumer(depth, deepest_state)
 
         return TransformerOutput(
             last_hidden_state=last_hidden_state,
@@ -962,6 +1018,7 @@ class TransformerStack(nn.Module):
             attentions=attentions,
             s_max=full_s_max,
             sae_hidden_states=sae_hidden_states,
+            captured_hidden_states=captured_hidden_states,
         )
 
     @torch.compiler.disable
@@ -975,6 +1032,7 @@ class TransformerStack(nn.Module):
         dtype: torch.dtype | None = None,
         effective_backend: AttentionBackend | None = None,
         output_attentions: bool = False,
+        validate_mask: bool = True,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None, BlockMask | None]:
         mask_name = "sequence_id" if sequence_id is not None else "attention_mask"
         mask_pattern = sequence_id if sequence_id is not None else attention_mask
@@ -1006,7 +1064,7 @@ class TransformerStack(nn.Module):
         attention_mask_2d = (
             mask_pattern if mask_pattern.dtype == torch.bool else mask_pattern != -1
         )
-        if not bool(attention_mask_2d.any(dim=1).all()):
+        if validate_mask and not bool(attention_mask_2d.any(dim=1).all()):
             raise ValueError("attention_mask must keep at least one valid key per batch row.")
 
         if mask_pattern.dtype == torch.bool:
@@ -1092,6 +1150,9 @@ class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel):
         "flash_attention_3",
     )
     _fastplms_attention_auto_order = ("sdpa",)
+    # embed_dataset(taps=...) runs _embed_taps; families without this marker reject taps.
+    embedding_tap_support: ClassVar[bool] = True
+    embedding_streaming_tap_support: ClassVar[bool] = True
 
     def __init__(self, config: ESMplusplusConfig, *args: object, **kwargs: object) -> None:
         super().__init__(config, *args, **kwargs)
@@ -1367,14 +1428,78 @@ class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel):
             if output.attentions is not None
             else None
         )
+        captured_hidden_states = (
+            {
+                index: state[:, :sequence_length]  # (b, l, d)
+                for index, state in output.captured_hidden_states.items()
+            }
+            if output.captured_hidden_states is not None
+            else None
+        )
         return TransformerOutput(
-            last_hidden_state=output.last_hidden_state[:, :sequence_length],
+            last_hidden_state=(
+                output.last_hidden_state[:, :sequence_length]
+                if output.last_hidden_state is not None
+                else None
+            ),
             hidden_states=hidden_states,
             attentions=attentions,
             s_max=output.s_max,
             sae_hidden_states=output.sae_hidden_states,
+            captured_hidden_states=captured_hidden_states,
         )
 
+    @property
+    def embedding_tap_state_count(self) -> int:
+        """Hidden states a tap can name: each block's input, then the final normalized state."""
+
+        return int(self.config.num_hidden_layers) + 1
+
+    def _embed_taps(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        layers: tuple[int, ...],
+        *,
+        stream_layers: tuple[int, ...] = (),
+        state_consumer: Callable[[int, torch.Tensor], None] | None = None,
+        assume_valid_mask: bool = False,
+    ) -> dict[int, torch.Tensor]:
+        """Hidden states at the FastPLMs indices ``layers`` from one pass.
+
+        The pass stops once the deepest requested state exists. It pads, enters, and trims the
+        FP8 context exactly as ``forward`` does, so each state equals the one
+        ``forward(output_hidden_states=True)`` returns for the same batch. ``assume_valid_mask``
+        skips the device-side mask check, which would stall the host on every batch.
+        """
+        # input_ids, attention_mask: (b, l)
+        if not layers and not stream_layers:
+            raise ValueError("_embed_taps needs at least one hidden-state index.")
+        input_ids, attention_mask, _, _, original_length = self._pad_fp8_inputs(
+            input_ids, attention_mask, None, None
+        )  # input_ids, attention_mask: (b, l_fp8), l_fp8 = l unless FP8 pads to 16
+        x = self.embed(input_ids)  # (b, l_fp8, d)
+
+        def consume(index: int, state: torch.Tensor) -> None:
+            # Reducer masks describe original tokens, never the alignment-only FP8 padding.
+            trimmed = state[:, :original_length] if original_length is not None else state
+            state_consumer(index, trimmed)  # (b, l, d)
+
+        with _esmplusplus_fp8_context(self._esmc_fp8, self.device):
+            output = self.transformer(
+                x=x,
+                attention_mask=attention_mask,
+                capture_layers=layers,
+                stop_after_layer=max((*layers, *stream_layers)),
+                stream_layers=stream_layers,
+                state_consumer=consume if state_consumer is not None else None,
+                assume_valid_mask=assume_valid_mask,
+            )
+        captured = self._trim_transformer_output(output, original_length).captured_hidden_states
+        if captured is None and layers:
+            raise RuntimeError("ESM++ did not capture the requested hidden states.")
+        return captured or {}  # each: (b, l, d); streaming-only runs retain no states
+
     @property
     def tokenizer(self) -> EsmSequenceTokenizer:
         """Construct the sequence tokenizer only when a raw-sequence API needs it."""

fastplms/models/esmfold2/embedding.py

@@ -110,7 +110,10 @@ class ESMFold2EmbeddingMixin:
     def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult:
         """Embed single-chain proteins using the learned 256-wide ESMFold2 summary."""
 
-        return embed_dataset(self, inputs, **kwargs)
+        embeddings = embed_dataset(self, inputs, **kwargs)
+        # ESMFold2 serves no one-pass taps, and embed_dataset refuses taps= before inference.
+        assert isinstance(embeddings, EmbeddingResult)
+        return embeddings
 
 
 __all__ = ["ESMFold2EmbeddingMixin"]

fastplms/models/esmfold2/esmfold2_conformers.py

@@ -116,6 +116,9 @@ def _open_verified_asset(
                 remaining -= len(chunk)
             copied_size = snapshot.tell()
             extra_byte = source.read(1)
+            # Release the asset once the snapshot holds its bytes. An open handle
+            # would block another process from replacing the path on Windows.
+            source.close()
             if remaining or extra_byte:
                 observed_size = copied_size if remaining else copied_size + len(extra_byte)
                 raise ValueError(

fastplms/models/esmfold2/esmfold2_molecular_complex.py

@@ -242,14 +242,21 @@ def _protein_sequence_and_indices(
     return protein_indices, "".join(sequence), chain_ids, entity_ids, sym_ids, confidences
 
 
-def _protein_entity_metadata_value(value: int | str) -> int | str:
-    """Restore the numeric entity labels used by ProteinComplex metadata."""
+def _protein_entity_metadata_value(value: int | str, entity: int) -> int:
+    """Restore the numeric entity label ProteinComplex metadata carries for `entity`.
+
+    A complex converted from a ProteinComplex keeps its labels as numeric strings. A folded
+    complex records the entity type ("polymer", "non-polymer") under the same key instead, which
+    is no label, so the entity keeps its own number. ProteinChain accepts only an integer label.
+    """
+    if isinstance(value, int) and not isinstance(value, bool):
+        return value
     if isinstance(value, str):
         try:
             return int(value)
         except ValueError:
             pass
-    return value
+    return entity
 
 
 def _atom37_from_flat(
@@ -787,7 +794,7 @@ class MolecularComplex:
         metadata = ProteinComplexMetadata(
             entity_lookup={
                 int(entity): _protein_entity_metadata_value(
-                    self.metadata.entity_lookup.get(int(entity), int(entity))
+                    self.metadata.entity_lookup.get(int(entity), int(entity)), int(entity)
                 )
                 for entity in unique_entities
             },

fastplms/models/esmfold2/esmfold2_parsing.py

@@ -10,9 +10,22 @@ from contextlib import AbstractContextManager, nullcontext
 from pathlib import Path
 from typing import NamedTuple, TextIO
 
+from ...embeddings.inputs import FastaDialect, scan_fasta_lines
 from .esmfold2_utils_types import PathOrBuffer
 
 
+# Lines are used as given and `#` lines are comments. A header keeps its full text, and sequence data
+# before the first header is skipped.
+ESMFOLD2_FASTA = FastaDialect(
+    strip_lines=False,
+    comment_prefix="#",
+    first_word_header=False,
+    squeeze_sequence_whitespace=False,
+    orphan_message=None,
+    empty_message="Found no sequences in input",
+)
+
+
 class FastaEntry(NamedTuple):
     """One FASTA record in source order."""
 
@@ -23,27 +36,8 @@ class FastaEntry(NamedTuple):
 def parse_fasta(text: str) -> Generator[FastaEntry, None, None]:
     """Yield records from FASTA text without normalizing sequence symbols."""
 
-    header: str | None = None
-    sequence_lines: list[str] = []
-    found_record = False
-
-    for line in text.splitlines():
-        if not line or line.startswith("#"):
-            continue
-        if line.startswith(">"):
-            if header is not None:
-                found_record = True
-                yield FastaEntry(header, "".join(sequence_lines))
-            header = line[1:].strip()
-            sequence_lines.clear()
-        elif header is not None:
-            sequence_lines.append(line)
-
-    if header is not None:
-        found_record = True
-        yield FastaEntry(header, "".join(sequence_lines))
-    if not found_record:
-        raise ValueError("Found no sequences in input")
+    for record in scan_fasta_lines(text.splitlines(), ESMFOLD2_FASTA, source="input"):
+        yield FastaEntry(record.header, record.sequence)
 
 
 def _open_reader(source: PathOrBuffer) -> AbstractContextManager[TextIO]:

fastplms/models/esmfold2/modeling_esmfold2.py

@@ -57,6 +57,7 @@ from .esmfold2_constants_esm3 import (
 )
 from .modeling_esmfold2_common import (
     CHAR_VOCAB_SIZE,
+    DEFAULT_MAX_INFERENCE_SIGMA,
     MAX_ATOMIC_NUMBER,
     MSA_CONDITIONING_INPUT_NAMES,
     NUM_RES_TYPES,
@@ -1490,7 +1491,7 @@ class ESMFold2Model(
         early_exit: bool = False,
         noise_scale: float | None = None,
         step_scale: float | None = None,
-        max_inference_sigma: float | None = None,
+        max_inference_sigma: float | None = DEFAULT_MAX_INFERENCE_SIGMA,
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
@@ -1793,7 +1794,7 @@ class ESMFold2Model(
         seed: int | None = None,
         noise_scale: float | None = None,
         step_scale: float | None = None,
-        max_inference_sigma: int | None = None,
+        max_inference_sigma: float | None = DEFAULT_MAX_INFERENCE_SIGMA,
         early_exit: bool = False,
         complex_id: str = "pred",
         verbose: bool = False,

fastplms/models/esmfold2/modeling_esmfold2_common.py

@@ -62,6 +62,9 @@ _VALID_BACKENDS = (None, BACKEND_FUSED, BACKEND_CUEQ)
 ATOM_ATTENTION_DENSE = "dense"
 ATOM_ATTENTION_WINDOWED = "windowed"
 _VALID_ATOM_ATTENTION = (ATOM_ATTENTION_DENSE, ATOM_ATTENTION_WINDOWED)
+# The official sampler starts every fold at this noise level unless a caller overrides it;
+# the public forwards and fold() default to it so an omitted argument keeps the cap.
+DEFAULT_MAX_INFERENCE_SIGMA = 256.0
 MSA_CONDITIONING_INPUT_NAMES = (
     "msa",
     "msa_attention_mask",
@@ -2793,7 +2796,7 @@ class DiffusionStructureHead(nn.Module):
         token_attention_mask: Tensor | None = None,
         num_diffusion_samples: int = 1,
         num_sampling_steps: int | None = None,
-        max_inference_sigma: float | None = 256.0,
+        max_inference_sigma: float | None = DEFAULT_MAX_INFERENCE_SIGMA,
         noise_scale: float | None = None,
         step_scale: float | None = None,
         return_atom_repr: bool = False,

fastplms/models/esmfold2/modeling_esmfold2_experimental.py

@@ -21,8 +21,8 @@ from tqdm.auto import tqdm
 from transformers.modeling_utils import PreTrainedModel
 
 from .attention import ESMFold2AttentionMixin
-from .configuration_esmfold2 import ESMFold2Config
 from .confidence_checkpoint import install_confidence_checkpoint
+from .configuration_esmfold2 import ESMFold2Config
 from .embedding import ESMFold2EmbeddingMixin
 from .modeling_esmfold2 import (
     ESMCPrecision,
@@ -38,6 +38,7 @@ from .modeling_esmfold2 import (
 )
 from .modeling_esmfold2_common import (
     CHAR_VOCAB_SIZE,
+    DEFAULT_MAX_INFERENCE_SIGMA,
     MAX_ATOMIC_NUMBER,
     MSA_CONDITIONING_INPUT_NAMES,
     NUM_RES_TYPES,
@@ -731,7 +732,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         provide_soft_sequence_to_msa_and_profile: bool = True,
         noise_scale: float | None = None,
         step_scale: float | None = None,
-        max_inference_sigma: float | None = None,
+        max_inference_sigma: float | None = DEFAULT_MAX_INFERENCE_SIGMA,
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
@@ -1094,7 +1095,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         seed: int | None = None,
         noise_scale: float | None = None,
         step_scale: float | None = None,
-        max_inference_sigma: int | None = None,
+        max_inference_sigma: float | None = DEFAULT_MAX_INFERENCE_SIGMA,
         early_exit: bool = False,
         complex_id: str = "pred",
         verbose: bool = False,

fastplms/registry.py

@@ -87,6 +87,8 @@ _ROOT_FIELDS = frozenset(
         "families",
         "models",
         "runtime_assets",
+        "golden_artifacts",
+        "sparse_autoencoders",
     }
 )
 _UPSTREAM_FIELDS = frozenset(
@@ -94,6 +96,7 @@ _UPSTREAM_FIELDS = frozenset(
         "id",
         "path",
         "url",
+        "fetch_url",
         "revision",
         "license",
         "license_files",
@@ -177,6 +180,17 @@ _RUNTIME_ASSET_FIELDS = frozenset(
         "offline_behavior",
     }
 )
+_GOLDEN_ARTIFACT_FIELDS = frozenset(
+    {
+        "id",
+        "repository",
+        "revision",
+        "path",
+        "sha256",
+        "size",
+        "offline_behavior",
+    }
+)
 
 
 class RegistryError(ValueError):
@@ -262,6 +276,20 @@ class CheckpointSource:
         return MappingProxyType({item.path: item for item in self.files})
 
 
+@dataclass(frozen=True, slots=True)
+class SparseAutoencoderSpec:
+    """Pinned released SAE and its input contract, not its unknown training revision."""
+
+    id: str
+    base_model: str
+    checkpoint: CheckpointSource
+    layer: int
+    input_width: int
+    k: int
+    codebook_dim: int
+    input_kind: Literal["hidden_state"] = "hidden_state"
+
+
 @dataclass(frozen=True, slots=True)
 class OracleAsset:
     """Hash-pinned external file required by a native parity oracle."""
@@ -297,6 +325,19 @@ class OfficialGolden:
     tensors: FileDigest
 
 
+@dataclass(frozen=True, slots=True)
+class GoldenArtifact:
+    """Durable Hub copy of one model's official golden tensors at an immutable revision."""
+
+    model_id: str
+    repository: str
+    revision: str
+    path: str
+    sha256: str
+    size: int
+    offline_behavior: str
+
+
 @dataclass(frozen=True, slots=True)
 class UpstreamSource:
     """Pinned official implementation used as a parity oracle."""
@@ -309,6 +350,13 @@ class UpstreamSource:
     license_files: tuple[str, ...]
     license_digests: tuple[FileDigest, ...] = ()
     distribution_files: tuple[FileDigest, ...] = ()
+    # Where the pinned revision is fetched when `url`, the source of record, no longer serves it.
+    fetch_url: str = ""
+
+    @property
+    def clone_url(self) -> str:
+        """The URL that `.gitmodules` and source archives use."""
+        return self.fetch_url or self.url
 
 
 @dataclass(frozen=True, slots=True)
@@ -478,6 +526,8 @@ class ModelRegistry(Mapping[str, ModelSpec]):
         runtime_assets: Mapping[str, RuntimeAsset] = MappingProxyType({}),
         attention_kernels: Mapping[str, AttentionKernelSpec] = MappingProxyType({}),
         legal_files: tuple[FileDigest, ...] = (),
+        golden_artifacts: Mapping[str, GoldenArtifact] = MappingProxyType({}),
+        sparse_autoencoders: Mapping[str, SparseAutoencoderSpec] = MappingProxyType({}),
     ) -> None:
         self.schema_version = schema_version
         self.upstreams = MappingProxyType(dict(upstreams))
@@ -486,6 +536,8 @@ class ModelRegistry(Mapping[str, ModelSpec]):
         self._models = MappingProxyType(dict(models))
         self.runtime_assets = MappingProxyType(dict(runtime_assets))
         self.legal_files = legal_files
+        self.golden_artifacts = MappingProxyType(dict(golden_artifacts))
+        self.sparse_autoencoders = MappingProxyType(dict(sparse_autoencoders))
 
     def __getitem__(self, key: str) -> ModelSpec:
         return self._models[key]
@@ -501,6 +553,22 @@ class ModelRegistry(Mapping[str, ModelSpec]):
             raise KeyError(family_id)
         return tuple(model for model in self._models.values() if model.family.id == family_id)
 
+    def sae_for_base(
+        self, base_model: str, *, layer: int, k: int, codebook_dim: int
+    ) -> SparseAutoencoderSpec:
+        """Resolve an exact registered SAE selection without a Hub lookup or fallback."""
+        matches = tuple(
+            spec for spec in self.sparse_autoencoders.values()
+            if (spec.base_model, spec.layer, spec.k, spec.codebook_dim)
+            == (base_model, layer, k, codebook_dim)
+        )
+        if len(matches) != 1:
+            raise RegistryError(
+                f"Expected one pinned SAE for {base_model}, layer={layer}, k={k}, "
+                f"codebook_dim={codebook_dim}; found {len(matches)}."
+            )
+        return matches[0]
+
     def supported_attention_dtypes(
         self,
         family_id: str,
@@ -988,8 +1056,11 @@ def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]:
             raise RegistryError(f"Duplicate upstream path: {path!r}")
         paths.add(path)
         url = _require_str(value, "url", context)
-        if not url.startswith("https://github.com/") or not url.endswith(".git"):
-            raise RegistryError(f"{context}.url must be an HTTPS GitHub clone URL.")
+        fetch_url = _optional_str(value, "fetch_url", context) or ""
+        for field, candidate in (("url", url), ("fetch_url", fetch_url)):
+            is_github = candidate.startswith("https://github.com/")
+            if candidate and not (is_github and candidate.endswith(".git")):
+                raise RegistryError(f"{context}.{field} must be an HTTPS GitHub clone URL.")
         license_files = _require_str_list(value, "license_files", context)
         license_digests = _require_digest_list(value, "license_digests", context)
         if tuple(item.path for item in license_digests) != license_files:
@@ -1026,6 +1097,7 @@ def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]:
             license_files=license_files,
             license_digests=license_digests,
             distribution_files=distribution_files,
+            fetch_url=fetch_url,
         )
     return upstreams
 
@@ -1288,6 +1360,67 @@ def _parse_runtime_assets(
     return runtime_assets
 
 
+def _parse_golden_artifacts(
+    raw: object,
+    models: Mapping[str, ModelSpec],
+) -> dict[str, GoldenArtifact]:
+    """Parse the optional Hub copies of official golden tensors, keyed by model ID.
+
+    A copy must be byte-identical to the tensors its model's official_golden pins, so its
+    SHA-256 must equal that digest and the existing integrity check covers both files.
+    """
+    if raw is None:
+        return {}
+    if not isinstance(raw, list):
+        raise RegistryError("golden_artifacts must be an array of tables.")
+    golden_artifacts: dict[str, GoldenArtifact] = {}
+    for index, value in enumerate(raw):
+        context = f"golden_artifacts[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{context} must be a table.")
+        _reject_unknown_fields(value, _GOLDEN_ARTIFACT_FIELDS, context)
+        model_id = _require_str(value, "id", context)
+        model = models.get(model_id)
+        if model is None:
+            raise RegistryError(f"{context}.id references unknown model {model_id!r}.")
+        if model.official_golden is None:
+            raise RegistryError(f"{context}.id names {model_id!r}, which has no official_golden.")
+        if model_id in golden_artifacts:
+            raise RegistryError(f"Duplicate golden artifact for model {model_id!r}.")
+        repository = _require_str(value, "repository", context)
+        if _REPOSITORY_ID_RE.fullmatch(repository) is None:
+            raise RegistryError(f"{context}.repository must be a Hugging Face repository ID.")
+        revision = _require_str(value, "revision", context)
+        _validate_revision(revision, f"{context}.revision")
+        path = _require_str(value, "path", context)
+        expected_path = f"goldens/{model_id}.safetensors"
+        if path != expected_path:
+            raise RegistryError(f"{context}.path must be {expected_path!r}.")
+        sha256 = _require_str(value, "sha256", context)
+        if sha256 != model.official_golden.tensors.digest:
+            raise RegistryError(
+                f"{context}.sha256 must equal the official_golden tensors digest of {model_id!r}."
+            )
+        size = value.get("size")
+        if isinstance(size, bool) or not isinstance(size, int) or size <= 0:
+            raise RegistryError(f"{context}.size must be a positive byte count.")
+        offline_behavior = _require_str(value, "offline_behavior", context)
+        if offline_behavior not in _ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS:
+            raise RegistryError(
+                f"{context}.offline_behavior is unsupported: {offline_behavior!r}."
+            )
+        golden_artifacts[model_id] = GoldenArtifact(
+            model_id=model_id,
+            repository=repository,
+            revision=revision,
+            path=path,
+            sha256=sha256,
+            size=size,
+            offline_behavior=offline_behavior,
+        )
+    return golden_artifacts
+
+
 def _parse_confidence_adaptation(
     table: Mapping[str, Any], context: str, model_id: str | None = None
 ) -> ConfidenceAdaptation | None:
@@ -1667,6 +1800,57 @@ def _validate_registry(
             )
 
 
+def _parse_sparse_autoencoders(
+    raw: object, models: Mapping[str, ModelSpec]
+) -> dict[str, SparseAutoencoderSpec]:
+    """Validate optional SAE records independently of the base-model mapping."""
+    if raw is None:
+        return {}
+    if not isinstance(raw, list):
+        raise RegistryError("sparse_autoencoders must be an array of tables.")
+    allowed = frozenset({
+        "id", "base_model", "layer", "input_width", "k", "codebook_dim", "input_kind",
+        "checkpoint_repo", "checkpoint_revision", "checkpoint_files",
+    })
+    records: dict[str, SparseAutoencoderSpec] = {}
+    selections: set[tuple[str, int, int, int]] = set()
+    for index, table in enumerate(raw):
+        context = f"sparse_autoencoders[{index}]"
+        if not isinstance(table, dict):
+            raise RegistryError(f"{context} must be a table.")
+        _reject_unknown_fields(table, allowed, context)
+        identifier = _require_str(table, "id", context)
+        if _IDENTIFIER_RE.fullmatch(identifier) is None or identifier in records:
+            raise RegistryError(f"{context} has an invalid or duplicate SAE id: {identifier!r}.")
+        base = _require_str(table, "base_model", context)
+        if base not in models or models[base].family.id != "esm_plusplus":
+            raise RegistryError(f"{context}.base_model must name a registered ESMC base.")
+        numbers: dict[str, int] = {}
+        for field in ("layer", "input_width", "k", "codebook_dim"):
+            value = table.get(field)
+            if type(value) is not int or value < (0 if field == "layer" else 1):
+                raise RegistryError(f"{context}.{field} must be a valid integer dimension.")
+            numbers[field] = value
+        if numbers["k"] > numbers["codebook_dim"]:
+            raise RegistryError(f"{context}.k exceeds codebook_dim.")
+        if _require_str(table, "input_kind", context) != "hidden_state":
+            raise RegistryError(f"{context} requires hidden_state input, not residual updates.")
+        checkpoint = _parse_checkpoint(table, "checkpoint", context)
+        expected_files = {"config.json", f"layer_{numbers['layer']}.safetensors"}
+        if set(checkpoint.file_map) != expected_files or any(
+            item.algorithm != "sha256" for item in checkpoint.files
+        ):
+            raise RegistryError(f"{context} must SHA-256 pin exactly {sorted(expected_files)}.")
+        selection = (base, numbers["layer"], numbers["k"], numbers["codebook_dim"])
+        if selection in selections:
+            raise RegistryError(f"{context} duplicates an SAE selection: {selection}.")
+        selections.add(selection)
+        records[identifier] = SparseAutoencoderSpec(
+            id=identifier, base_model=base, checkpoint=checkpoint, **numbers
+        )
+    return records
+
+
 def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry:
     try:
         manifest = tomllib.loads(raw_bytes.decode("utf-8"))
@@ -1685,6 +1869,8 @@ def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry:
     runtime_assets = _parse_runtime_assets(manifest.get("runtime_assets"), families)
     models = _parse_models(manifest.get("models"), families)
     _validate_registry(upstreams, attention_kernels, families, models)
+    golden_artifacts = _parse_golden_artifacts(manifest.get("golden_artifacts"), models)
+    sparse_autoencoders = _parse_sparse_autoencoders(manifest.get("sparse_autoencoders"), models)
     return ModelRegistry(
         schema_version=1,
         upstreams=upstreams,
@@ -1693,6 +1879,8 @@ def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry:
         models=models,
         runtime_assets=runtime_assets,
         legal_files=legal_files,
+        golden_artifacts=golden_artifacts,
+        sparse_autoencoders=sparse_autoencoders,
     )
 
 
@@ -1729,6 +1917,7 @@ __all__ = [
     "CheckpointSource",
     "FileDigest",
     "GenerationContract",
+    "GoldenArtifact",
     "ModelFamily",
     "ModelRegistry",
     "ModelSpec",
@@ -1737,6 +1926,7 @@ __all__ = [
     "RuntimeAsset",
     "RuntimeAssetTrustKind",
     "RuntimeExtra",
+    "SparseAutoencoderSpec",
     "TestTier",
     "UpstreamSource",
     "VramTier",

modeling_fastplms.py

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

requirements.txt

@@ -1,20 +1,20 @@
 # Direct runtime dependencies for Synthyra/ESMFold2-Fast.
 # FastPLMs source is embedded in this model repository.
-torch>=2.13,<2.14
-transformers>=5.13,<5.14
-huggingface-hub>=0.34,<2
-tokenizers>=0.22,<0.23
-safetensors>=0.5,<1
-numpy>=1.26,<3
-einops>=0.8,<1
-tqdm>=4.67,<5
-accelerate>=1.10,<2
-biopython>=1.85,<2
-biotite>=1.4,<2
-brotli>=1.1,<2
-msgpack>=1.1,<2
-msgpack-numpy>=0.4.8,<1
-omegaconf>=2.3,<3
-rdkit>=2025.9,<2027
-scipy>=1.15,<2
-zstandard>=0.23,<1
+torch>=2.14
+transformers>=5.17
+huggingface-hub>=1.32
+tokenizers>=0.23
+safetensors>=0.8
+numpy>=2.5
+einops>=0.8
+tqdm>=4.70
+accelerate>=1.15
+biopython>=1.88
+biotite>=1.7
+brotli>=1.2
+msgpack>=1.2
+msgpack-numpy>=0.4.8
+omegaconf>=2.3
+rdkit>=2026.3
+scipy>=1.18
+zstandard>=0.25