Update FastPLMs files

LICENSES/protein-ttt/SOURCE_RECORD.md

@@ -0,0 +1,8 @@
+# ProteinTTT provenance
+
+FastPLMs uses `anton-bushuiev/ProteinTTT` revision
+`fde2817cd84b936167cc76ccabf31e5c0fe49962` as the official reference for the
+optional protein test-time training workflow. The repository is pinned at
+`vendor/upstream/protein-ttt/` and is not a production dependency or runtime
+image component. The accompanying `LICENSE` is the verbatim MIT text from that
+revision.

README.md

@@ -10,7 +10,7 @@ tags:
 
 # Synthyra/ESMFold2-Fast
 
-This checkpoint packages the FastPLMs `ESMFold2` implementation.
+This checkpoint contains the FastPLMs `ESMFold2` implementation.
 
 Accepted inputs are raw amino-acid sequences or typed molecular-complex
 specifications; low-level forward accepts prepared feature tensors.
@@ -28,9 +28,7 @@ Supported Transformers entry points are `AutoConfig`, `AutoModel`.
 | Attention variants | Supported: `eager`, `sdpa`, `flex_attention` |
 | Compliance | Declared: exact release evidence is required |
 
-A supported interface is not a pretrained downstream predictor. Classification
-heads start untrained, and declared compliance metadata is not a claim that an
-arbitrary local build passed its release gate.
+A supported interface is not a pretrained downstream predictor. Classification heads start untrained. Compliance metadata does not show that a local build passed its release gate.
 
 ## Install and platform requirements
 
@@ -41,12 +39,12 @@ python -m pip install -r \
   "https://huggingface.co/Synthyra/ESMFold2-Fast/resolve/main/requirements.txt"
 ```
 
-The FastPLMs implementation itself is embedded in the model repository and loaded
-by Transformers through `trust_remote_code=True`.
+The FastPLMs implementation itself is embedded in the model repository.
+Transformers loads it through `trust_remote_code=True`.
 
-Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The artifact requirements include the direct structure dependencies. The published execution contract requires a CUDA device. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. The Hub quick start below requires network
-access on first download. For an air-gapped run, first build the manifest-pinned
-local artifact and use the offline form shown in the example.
+This model requires Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13. The artifact requirements include the structure dependencies. The release contract requires a CUDA device. The current validated target is the exact NVIDIA GH200 on Linux aarch64. Linux x86-64, CPU-only, Windows, and macOS structure runs are not release evidence. The Hub quick start needs network access for
+the first download. For an air-gapped run, build the manifest-pinned local
+artifact first and use the offline example.
 
 ## Quick start
 
@@ -62,21 +60,21 @@ model = AutoModel.from_pretrained(
 ```
 
 For offline validation, replace `model_id` with the manifest-built
-`dist/hub/ESMFold2-Fast` path and pass `local_files_only=True`.
+`dist/hub/ESMFold2-Fast` path. Pass `local_files_only=True`.
 
 ## Attention and compliance
 
 The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`. An unavailable
-requested backend raises instead of silently switching implementations.
-`output_attentions=True` may use the documented, one-call eager fallback solely
-to materialize attention tensors; the configured backend remains unchanged.
+requested backend raises. It does not silently change implementation.
+`output_attentions=True` can use the documented one-call eager fallback to
+materialize attention tensors. The configured backend does not change.
 
-This family declares the `compliance` tier. Release evidence binds the exact
+This family declares the `compliance` tier. Release evidence identifies the
 checkpoint, backend, dtype, hardware, inputs, and reference revision.
 
 ## PEFT fine-tuning
 
-Install the direct training dependencies, then attach LoRA to the loaded checkpoint:
+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"
@@ -95,22 +93,21 @@ peft_model = get_peft_model(
 )
 ```
 
-This checkpoint has no advertised classifier. Supply the task-specific
-objective and preserve any new head through `modules_to_save`.
+This checkpoint has no advertised classifier. Supply the task objective and
+preserve any new head through `modules_to_save`.
 All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and
-can be adapted with PEFT. The ESM2-specific shipped CLI is an example, not a
+can use PEFT. The ESM2-specific shipped CLI is an example, not a
 support boundary. Record the target modules, base revision, data identity, and
 trainable parameter scope.
 
 ## Alignment-conditioning contract
 
-This 24-block Fast checkpoint is inference-optimized for single-sequence
-conditioning and was trained without MSA conditioning. It is not
-MSA-conditioned and rejects `ProteinInput.msa` and low-level MSA-derived
-features. Typed multichain and multimolecule inputs remain supported when
-every protein chain uses `msa=None`. Use the corresponding full ESMFold2
-checkpoint for MSA-conditioned inference. This follows the official Biohub
-architecture description in [Appendix A.2.1](https://biohub.ai/papers/esm_protein.pdf).
+This 24-block Fast checkpoint is optimized for single-sequence inference. It
+was trained without MSA conditioning. It rejects `ProteinInput.msa` and low-level
+MSA-derived features. Typed multichain and multimolecule inputs remain supported
+when every protein chain uses `msa=None`. Use the full ESMFold2 checkpoint for
+MSA-conditioned inference. This follows the official Biohub architecture
+description in [Appendix A.2.1](https://biohub.ai/papers/esm_protein.pdf).
 
 
 ## Protein folding
@@ -152,15 +149,13 @@ complex_result = model.fold(
 print(complex_result.ptm, complex_result.plddt.mean().item())
 ```
 
-The typed interface also supports RNA, modifications, and covalent bonds.
-Protein MSA inputs are not supported by this Fast checkpoint; every protein
-chain must use `msa=None`. The public schema recognizes `PocketConditioning` and
-`DistogramConditioning`, but the pinned official forward consumes neither. Its
-feature builder hard-codes a zero pocket feature and constructs distogram tensors
-that the released model ignores. FastPLMs therefore rejects non-null pocket and
-distogram conditioning instead of silently ignoring scientific inputs. Prepared
-`ref_pos` values are component reference geometries created during featurization,
-not target coordinates.
+The typed interface also supports RNA, modifications, covalent bonds, and
+distogram conditioning. Protein MSA inputs are not supported by this Fast
+checkpoint; every protein chain must use `msa=None`. The public schema recognizes
+`PocketConditioning`, but the pinned official runtime discards it and hard-codes
+a zero pocket feature. FastPLMs therefore rejects non-null pocket conditioning
+instead of silently ignoring it. Prepared `ref_pos` values are component
+reference geometries created during featurization, not target coordinates.
 Predicted coordinates and confidence scores are outputs and do not establish
 biochemical activity.
 
@@ -244,9 +239,9 @@ print(adapted.ttt_metrics)
 ```
 
 Entering a gradient-enabled path reloads canonical BF16 ESMC weights. TTT adds
-latency and memory, can worsen a prediction, and does not calibrate confidence
-or establish biological validity. Folding TTT is result-scoped: its transient
-ESMC adapter modules are excluded from checkpoint state, so it is not a generic
+latency and memory and can worsen a prediction. It does not calibrate confidence
+or show biological validity. Folding TTT is result-scoped. Its transient ESMC
+adapter modules are excluded from checkpoint state. It is not a generic
 `save_pretrained` adapter-persistence path.
 
 ## Runtime contract
@@ -267,8 +262,8 @@ ESMC adapter modules are excluded from checkpoint state, so it is not a generic
 ## Release record
 
 - FastPLMs weights: `Synthyra/ESMFold2-Fast`
-- Runtime revision: recorded separately in the built artifact and published commit
-- Source-tree and runtime-bundle SHA-256: recorded in `provenance.json`
+- Runtime revision: recorded in the built artifact and published commit
+- Source-tree and runtime-bundle SHA-256: recorded in the source record
 - Official checkpoint: `biohub/ESMFold2-Fast`
 - Artifact source: `fast`
 - State transform: `identity`
@@ -276,19 +271,17 @@ ESMC adapter modules are excluded from checkpoint state, so it is not a generic
 - Release tiers: `check`, `compliance`, `structure`, `feature`, `artifact`, `benchmark`
 - Unresolved required file identities: `0`
 
-`provenance.json` records exact file identities, conversion, source revisions,
-legal texts, schema, and attestations. A nonzero unresolved count blocks release.
+The source record records exact file identities, conversion, source revisions,
+legal texts, schema, and attestations. A nonzero unresolved count blocks a release.
 
 ## Validation boundary
 
-Declared tiers compare applicable configuration, tokenizer behavior, state,
-and representative inference with the pinned reference. Metadata alone does
-not claim a build passed, a backend is faster, or an output is biologically
-valid.
+Declared tiers compare configuration, tokenizer behavior, state, and
+representative inference with the pinned reference. Metadata does not show that
+a build passed, that a backend is faster, or that an output is biologically valid.
 
 ## License
 
 Checkpoint terms: MIT. The Hub model-card identifier is
-`mit`. Applicable source licenses, notices, attribution,
-and conversion records are distributed with the local artifact. Review them
-before use.
+`mit`. The local artifact contains applicable source
+licenses, notices, attribution, and conversion records. Review them before use.

THIRD_PARTY_NOTICES.md

@@ -46,7 +46,7 @@ explicitly defines the repository release as including pretrained DPLM1 and
 DPLM2 weights, and the same revision carries the complete
 [Apache-2.0 license](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/LICENSE).
 FastPLMs records both checkpoint families as Apache-2.0 and distributes the
-verbatim license plus `LICENSES/dplm/PROVENANCE.md`. Converted weights retain
+verbatim license plus `LICENSES/dplm/SOURCE_RECORD.md`. Converted weights retain
 those terms and remain subject to the ordinary artifact and publication gates.
 
 ## Biohub
@@ -80,7 +80,7 @@ TorchMetrics, Lightning Utilities, and NVIDIA DLLogger. Their exact versions or
 revision are pinned in `docker/constraints/esmfold.txt`; OpenFold imports them
 eagerly, and FastPLMs production code does not depend on them. DLLogger's exact
 source identity and installed-license handling are recorded in
-`LICENSES/dllogger/PROVENANCE.md`.
+`LICENSES/dllogger/SOURCE_RECORD.md`.
 
 ## ProteinTTT
 
@@ -93,7 +93,7 @@ revision-specific provenance are under `LICENSES/protein-ttt/`.
 For every supported family, `src/fastplms/models.toml` records an immutable
 official checkpoint revision, an immutable FastPLMs checkpoint revision, file
 digests, a named state transformation, and a mechanism-level conversion record.
-Generated artifacts reproduce that record in `provenance.json`. A release or
+Generated artifacts reproduce that record in `source-record.json`. A release or
 artifact build must fail when a required file identity, legal text, attribution
 notice, modified-file notice, upstream revision, or conversion record is absent
 or differs from its manifest digest.

fastplms/models.toml

@@ -88,7 +88,7 @@ license_files = ["LICENSE"]
 license_digests = ["LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"]
 distribution_files = [
   "LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
-  "PROVENANCE.md=sha256:a659f74be9073cf1ad2d2f7071531ca56959b421f111152cf4c41184ace5970e",
+  "SOURCE_RECORD.md=sha256:a659f74be9073cf1ad2d2f7071531ca56959b421f111152cf4c41184ace5970e",
 ]
 
 [[upstreams]]
@@ -122,7 +122,7 @@ license_files = ["LICENSE"]
 license_digests = ["LICENSE=sha256:da6d3703ed11cbe42bd212c725957c98da23cbff1998c05fa4b3d976d1a58e93"]
 distribution_files = [
   "LICENSE=sha256:da6d3703ed11cbe42bd212c725957c98da23cbff1998c05fa4b3d976d1a58e93",
-  "PROVENANCE.md=sha256:950adb94daf15e646ddf226dacfe2a8e77801aa0793e439a9a3490a48eb666e7",
+  "SOURCE_RECORD.md=sha256:950adb94daf15e646ddf226dacfe2a8e77801aa0793e439a9a3490a48eb666e7",
 ]
 
 [[upstreams]]
@@ -136,7 +136,7 @@ license_digests = ["LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c1
 distribution_files = [
   "LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
   "MODIFICATIONS.md=sha256:fd6f0aa1086a0c996cf967b326d18e965660cda0ad5c7f36a3474a8490720da3",
-  "PROVENANCE.md=sha256:48c903db43a217a3126afaefbac60b7ddac7efda2dfcc0cbff0bffc7d6c30081",
+  "SOURCE_RECORD.md=sha256:48c903db43a217a3126afaefbac60b7ddac7efda2dfcc0cbff0bffc7d6c30081",
 ]
 
 [[upstreams]]
@@ -149,7 +149,7 @@ license_files = ["LICENSE"]
 license_digests = ["LICENSE=sha256:bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df"]
 distribution_files = [
   "LICENSE=sha256:bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df",
-  "PROVENANCE.md=sha256:dc641c37353c2efd50ccbdb316ca4aae495ec02c1563e0e15bac92f75fc482e5",
+  "SOURCE_RECORD.md=sha256:dc641c37353c2efd50ccbdb316ca4aae495ec02c1563e0e15bac92f75fc482e5",
 ]
 
 [families.esm2]
@@ -187,7 +187,8 @@ reference_adapter = "tests.parity.support.reference_adapters.esm_plusplus"
 attention = ["eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"]
 dtypes = ["float32", "bfloat16"]
 bf16_execution = "static_parameters"
-precisions = ["default"]
+precisions = ["default", "fp8"]
+experimental_precisions = ["fp8"]
 vram_tier = "sequence"
 checkpoint_license = "MIT"
 hub_license = "mit"
@@ -267,7 +268,7 @@ checkpoint_license = "Apache-2.0"
 hub_license = "apache-2.0"
 weights_publication_allowed = true
 state_transform = "dplm_to_fastplms_v1"
-conversion_provenance = "Input: the pinned official DPLM1 checkpoint. Transformation: apply dplm_to_fastplms_v1, omitting the unused absolute-position table for rotary checkpoints and materializing the tied input/output embedding values as independent tensors. Output: the pinned Synthyra DPLM checkpoint. Validation: release parity compares exact state identity after the declared transform, tokenizer behavior, generation, and inference. License basis: the pinned ByteDance DPLM Apache-2.0 LICENSE and README explicitly scope the repository release to the pretrained DPLM1 and DPLM2 weights; immutable evidence is recorded in LICENSES/dplm/PROVENANCE.md. Limitation: redistribution remains subject to Apache-2.0 and the pinned source record; no broader rights are inferred."
+conversion_provenance = "Input: the pinned official DPLM1 checkpoint. Transformation: apply dplm_to_fastplms_v1, omitting the unused absolute-position table for rotary checkpoints and materializing the tied input/output embedding values as independent tensors. Output: the pinned Synthyra DPLM checkpoint. Validation: release parity compares exact state identity after the declared transform, tokenizer behavior, generation, and inference. License basis: the pinned ByteDance DPLM Apache-2.0 LICENSE and README explicitly scope the repository release to the pretrained DPLM1 and DPLM2 weights; immutable evidence is recorded in LICENSES/dplm/SOURCE_RECORD.md. Limitation: redistribution remains subject to Apache-2.0 and the pinned source record; no broader rights are inferred."
 representative = "dplm_150m"
 documentation = "docs/models.md#dplm"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
@@ -291,7 +292,7 @@ checkpoint_license = "Apache-2.0"
 hub_license = "apache-2.0"
 weights_publication_allowed = true
 state_transform = "dplm2_to_fastplms_v1"
-conversion_provenance = "Input: the pinned official DPLM2 checkpoint. Transformation: apply dplm2_to_fastplms_v1, retaining the independent language-model head and trained encoder contact head while omitting the unused absolute-position table for rotary checkpoints. Output: the pinned Synthyra DPLM2 checkpoint. Validation: release parity compares exact keys and values after the declared omission, non-aliasing, tokenizer behavior, generation, and inference. License basis: the pinned ByteDance DPLM Apache-2.0 LICENSE and README explicitly scope the repository release to the pretrained DPLM1 and DPLM2 weights; immutable evidence is recorded in LICENSES/dplm/PROVENANCE.md. Limitation: no head exception is permitted by this source record, and redistribution remains subject to Apache-2.0."
+conversion_provenance = "Input: the pinned official DPLM2 checkpoint. Transformation: apply dplm2_to_fastplms_v1, retaining the independent language-model head and trained encoder contact head while omitting the unused absolute-position table for rotary checkpoints. Output: the pinned Synthyra DPLM2 checkpoint. Validation: release parity compares exact keys and values after the declared omission, non-aliasing, tokenizer behavior, generation, and inference. License basis: the pinned ByteDance DPLM Apache-2.0 LICENSE and README explicitly scope the repository release to the pretrained DPLM1 and DPLM2 weights; immutable evidence is recorded in LICENSES/dplm/SOURCE_RECORD.md. Limitation: no head exception is permitted by this source record, and redistribution remains subject to Apache-2.0."
 representative = "dplm2_150m"
 documentation = "docs/models.md#dplm2"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]

fastplms/models/esm_plusplus/modeling_esm_plusplus.py

@@ -2,13 +2,17 @@
 
 from __future__ import annotations
 
+import importlib
+import importlib.metadata
 import math
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass
+from functools import partial
+from typing import Any, ClassVar
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from dataclasses import dataclass
-from functools import partial
-from typing import ClassVar
 from einops import rearrange
 from tokenizers import Tokenizer
 from tokenizers.models import BPE
@@ -61,6 +65,151 @@ except ModuleNotFoundError as error:
     # Legacy flat Hub composites define every shared symbol above this block.
 
 
+_ESMC_FP8_LINEAR_SUFFIX = ".attn.out_proj"
+_ESMC_FP8_ALIGNMENT = 16
+
+
+@dataclass(frozen=True, slots=True)
+class ESMplusplusFP8Status:
+    """Resolved state of the explicit Transformer Engine ESMC FP8 path."""
+
+    enabled: bool
+    reason: str
+    device: str
+    transformer_engine_version: str | None
+    converted_projections: int
+
+    def as_dict(self) -> dict[str, str | int | bool | None]:
+        return asdict(self)
+
+
+def _transformer_engine_version() -> str | None:
+    try:
+        return importlib.metadata.version("transformer-engine")
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def _load_transformer_engine() -> tuple[Any, Any]:
+    """Load Transformer Engine lazily so ordinary ESM++ imports stay portable."""
+
+    try:
+        te = importlib.import_module("transformer_engine.pytorch")
+        recipe = importlib.import_module("transformer_engine.common.recipe")
+    except (ImportError, OSError, RuntimeError) as error:
+        raise RuntimeError(
+            f"Transformer Engine could not be imported: {type(error).__name__}: {error}"
+        ) from error
+    if not hasattr(recipe, "Float8CurrentScaling"):
+        raise RuntimeError(
+            "Transformer Engine does not expose Float8CurrentScaling, which is "
+            "required by the validated ESMC FP8 path."
+        )
+    return te, recipe
+
+
+def _te_fp8_capability(device: torch.device) -> tuple[bool, str]:
+    """Return whether the strict Transformer Engine FP8 path can run."""
+
+    if device.type != "cuda":
+        return False, "FP8 requires ESM++ on a CUDA device."
+    if not torch.cuda.is_available():
+        return False, "CUDA is unavailable."
+    try:
+        major, minor = torch.cuda.get_device_capability(device)
+    except (AssertionError, RuntimeError, ValueError) as error:
+        return False, f"CUDA capability query failed: {error}"
+    if not (major >= 9 or (major == 8 and minor >= 9)):
+        return False, f"CUDA capability {major}.{minor} does not support FP8."
+    try:
+        te, _ = _load_transformer_engine()
+    except RuntimeError as error:
+        return False, str(error)
+
+    probe = getattr(te, "is_fp8_available", None)
+    if probe is None:
+        try:
+            probe = importlib.import_module("transformer_engine.pytorch.fp8").is_fp8_available
+        except (ImportError, AttributeError, OSError, RuntimeError) as error:
+            return False, f"Transformer Engine has no usable FP8 probe: {error}"
+    try:
+        try:
+            result = probe(return_reason=True)
+        except TypeError:
+            result = probe()
+    except (OSError, RuntimeError) as error:
+        return False, f"Transformer Engine FP8 probe failed: {error}"
+    if isinstance(result, tuple):
+        available = bool(result[0])
+        detail = str(result[1]) if len(result) > 1 and result[1] else ""
+    else:
+        available = bool(result)
+        detail = ""
+    if not available:
+        return False, detail or "Transformer Engine reports FP8 unavailable."
+    return True, "Transformer Engine reports FP8 availability."
+
+
+def _convert_esmc_attention_outputs_to_te(
+    module: nn.Module,
+    *,
+    expected_projections: int,
+) -> tuple[str, ...]:
+    """Replace exactly one attention output projection per ESMC block."""
+
+    targets = [
+        (path, child)
+        for path, child in module.named_modules()
+        if isinstance(child, nn.Linear) and path.endswith(_ESMC_FP8_LINEAR_SUFFIX)
+    ]
+    if len(targets) != expected_projections:
+        raise RuntimeError(
+            "ESMC FP8 conversion expected exactly "
+            f"{expected_projections} attention output projections, found {len(targets)}."
+        )
+
+    te, _ = _load_transformer_engine()
+    modules = dict(module.named_modules())
+    converted: list[str] = []
+    for path, child in targets:
+        owner_path, name = path.rsplit(".", 1)
+        owner = modules[owner_path]
+        replacement = te.Linear(
+            child.in_features,
+            child.out_features,
+            bias=child.bias is not None,
+            params_dtype=child.weight.dtype,
+            device=child.weight.device,
+        )
+        with torch.no_grad():
+            replacement.weight.copy_(child.weight)
+            if child.bias is not None:
+                replacement.bias.copy_(child.bias)
+        replacement.eval().requires_grad_(False)
+        setattr(owner, name, replacement)
+        converted.append(path)
+    return tuple(converted)
+
+
+@contextmanager
+def _esmplusplus_fp8_context(enabled: bool, device: torch.device):
+    """Enter the validated BF16-storage, Transformer Engine FP8 context."""
+
+    if not enabled:
+        yield
+        return
+    if torch.is_grad_enabled():
+        raise RuntimeError("ESM++ FP8 is inference-only; use torch.inference_mode() or no_grad().")
+    te, recipe = _load_transformer_engine()
+    fp8_recipe = recipe.Float8CurrentScaling(
+        use_power_2_scales=False,
+        fp8_format=recipe.Format.HYBRID,
+    )
+    with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
+        with te.autocast(enabled=True, recipe=fp8_recipe):
+            yield
+
+
 class ESMplusplusConfig(PretrainedConfig):
     """Configuration class for ESM++ model.
 
@@ -616,6 +765,8 @@ class TransformerOutput(ModelOutput):
     hidden_states: tuple[torch.Tensor] | None = None
     attentions: tuple[torch.Tensor] | None = None
     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
 
 
 @dataclass
@@ -624,6 +775,7 @@ class ESMplusplusOutput(MaskedLMOutput):
 
     s_max: tuple[list[torch.Tensor], ...] | None = None
     last_hidden_state: torch.Tensor | None = None
+    sae_outputs: dict[str, torch.Tensor] | None = None
 
 
 @dataclass
@@ -631,6 +783,7 @@ class ESMplusplusSequenceClassifierOutput(SequenceClassifierOutput):
     """Sequence-classification output with optional attention diagnostics."""
 
     s_max: tuple[list[torch.Tensor], ...] | None = None
+    sae_outputs: dict[str, torch.Tensor] | None = None
 
 
 @dataclass
@@ -638,6 +791,7 @@ class ESMplusplusTokenClassifierOutput(TokenClassifierOutput):
     """Token-classification output with optional attention diagnostics."""
 
     s_max: tuple[list[torch.Tensor], ...] | None = None
+    sae_outputs: dict[str, torch.Tensor] | None = None
 
 
 class TransformerStack(nn.Module):
@@ -688,63 +842,30 @@ class TransformerStack(nn.Module):
         output_attentions: bool | None = False,
         output_s_max: bool | None = False,
         esmfold2_hidden_states: bool = False,
+        sae_layers: tuple[int, ...] = (),
     ) -> TransformerOutput:
         # x: (b, l, d); attention_mask, sequence_id: (b, l)
         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
         # 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.
-        if sequence_id is None and attention_mask is not None:
-            expected_shape = (x.shape[0], x.shape[1])
-            if attention_mask.ndim != 2 or tuple(attention_mask.shape) != expected_shape:
-                raise ValueError(
-                    f"attention_mask must have shape {expected_shape}; "
-                    f"received {tuple(attention_mask.shape)}."
-                )
-            attention_mask = attention_mask.to(device=x.device, dtype=torch.bool)
-            if not bool(attention_mask.any(dim=1).all()):
-                raise ValueError("attention_mask must keep at least one valid key per batch row.")
-        effective_backend = resolve_attention_backend_for_call(
-            self.attention_backend,
-            output_attentions=bool(output_attentions),
-        )
-
-        if sequence_id is None and attention_mask is not None:
-            attention_mask_2d, attention_mask_4d, flex_block_mask = (
-                self._sequence_id_attention_masks(
-                    sequence_id=attention_mask.to(device=x.device, dtype=torch.bool),
-                    batch_size=x.shape[0],
-                    seq_len=x.shape[1],
-                    device=x.device,
-                    dtype=x.dtype,
-                    effective_backend=effective_backend,
-                )
-            )
-        elif sequence_id is None:
-            attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask(
-                effective_backend=effective_backend,
+        attention_mask_2d, attention_mask_4d, flex_block_mask = (
+            self._prepare_attention_masks(
+                attention_mask=attention_mask,
+                sequence_id=sequence_id,
                 batch_size=x.shape[0],
                 seq_len=x.shape[1],
                 device=x.device,
-                attention_mask=attention_mask,
                 dtype=x.dtype,
-                mask_semantics="padding",
-            )
-        else:
-            attention_mask_2d, attention_mask_4d, flex_block_mask = (
-                self._sequence_id_attention_masks(
-                    sequence_id=sequence_id,
-                    batch_size=x.shape[0],
-                    seq_len=x.shape[1],
-                    device=x.device,
-                    dtype=x.dtype,
-                    effective_backend=effective_backend,
-                )
+                output_attentions=bool(output_attentions),
             )
+        )
 
-        for block in self.blocks:
+        for layer_index, block in enumerate(self.blocks):
             if output_hidden_states:
                 if hidden_states is None:
                     raise RuntimeError(
@@ -754,6 +875,8 @@ class TransformerStack(nn.Module):
                 # by the final normalized state. This gives n_layers + 1 states
                 # and, for ESMC-6B, the 81-state order consumed by ESMFold2.
                 hidden_states += (x,)
+            if sae_hidden_states is not None and layer_index in sae_layer_set:
+                sae_hidden_states[layer_index] = x
             if self.gradient_checkpointing and self.training:
                 x, attn_weights, s_max = self._gradient_checkpointing_func(
                     block.__call__,
@@ -782,52 +905,85 @@ class TransformerStack(nn.Module):
         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
 
         return TransformerOutput(
             last_hidden_state=last_hidden_state,
             hidden_states=hidden_states,
             attentions=attentions,
             s_max=full_s_max,
+            sae_hidden_states=sae_hidden_states,
         )
 
-    def _sequence_id_attention_masks(
+    @torch.compiler.disable
+    def _prepare_attention_masks(
         self,
-        sequence_id: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        sequence_id: torch.Tensor | None,
         batch_size: int,
         seq_len: int,
         device: torch.device,
         dtype: torch.dtype | None = None,
         effective_backend: AttentionBackend | None = None,
+        output_attentions: bool = False,
     ) -> 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
+        if mask_pattern is None:
+            backend = resolve_attention_backend_for_call(
+                self.attention_backend,
+                output_attentions=output_attentions,
+            )
+            return get_attention_mask(
+                effective_backend=backend,
+                batch_size=batch_size,
+                seq_len=seq_len,
+                device=device,
+                attention_mask=None,
+                dtype=dtype,
+                mask_semantics="padding",
+            )
+
         expected_shape = (batch_size, seq_len)
-        if sequence_id.ndim != 2 or tuple(sequence_id.shape) != expected_shape:
+        if mask_pattern.ndim != 2 or tuple(mask_pattern.shape) != expected_shape:
             raise ValueError(
-                f"sequence_id must have shape {expected_shape}; "
-                f"received {tuple(sequence_id.shape)}."
+                f"{mask_name} must have shape {expected_shape}; "
+                f"received {tuple(mask_pattern.shape)}."
             )
-        if sequence_id.device != device:
-            sequence_id = sequence_id.to(device=device)
-        backend = (
-            self.attention_backend
-            if effective_backend is None
-            else resolve_attention_backend(effective_backend)
+        if mask_pattern.device != device:
+            mask_pattern = mask_pattern.to(device=device)
+        if sequence_id is None:
+            mask_pattern = mask_pattern.to(dtype=torch.bool)
+        attention_mask_2d = (
+            mask_pattern if mask_pattern.dtype == torch.bool else mask_pattern != -1
         )
-        if sequence_id.dtype == torch.bool:
-            attention_mask_2d = sequence_id
+        if 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:
             # Biohub's boolean single-chain form groups biological positions
             # together and padding positions together. Padding queries remain
             # finite without allowing their states to enter residue attention.
-            attention_mask_4d = sequence_id[:, None, :, None] == sequence_id[:, None, None, :]
+            attention_mask_4d = (
+                mask_pattern[:, None, :, None] == mask_pattern[:, None, None, :]
+            )
         else:
-            attention_mask_2d = sequence_id != -1
-            attention_mask_4d = (sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2)).unsqueeze(
-                1
+            attention_mask_4d = (
+                mask_pattern.unsqueeze(-1) == mask_pattern.unsqueeze(-2)
+            ).unsqueeze(1)
+        backend = (
+            resolve_attention_backend_for_call(
+                self.attention_backend,
+                output_attentions=output_attentions,
             )
-        if not bool(attention_mask_2d.any(dim=1).all()):
-            raise ValueError("attention_mask must keep at least one valid key per batch row.")
+            if effective_backend is None
+            else resolve_attention_backend(effective_backend)
+        )
 
         if backend.is_flash:
-            if sequence_id.dtype != torch.bool:
+            if mask_pattern.dtype != torch.bool:
                 raise ValueError(
                     "ESM++ FlashAttention only supports boolean sequence_id padding masks. "
                     "Use eager, sdpa, or flex_attention for chain-aware integer sequence_id "
@@ -836,22 +992,22 @@ class TransformerStack(nn.Module):
             return attention_mask_2d, attention_mask_4d, None
 
         if backend == AttentionBackend.FLEX:
-            if sequence_id.dtype == torch.bool:
+            if mask_pattern.dtype == torch.bool:
 
                 def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
                     del head_idx
-                    return sequence_id[batch_idx, q_idx] == sequence_id[batch_idx, kv_idx]
+                    return mask_pattern[batch_idx, q_idx] == mask_pattern[batch_idx, kv_idx]
 
             else:
 
                 def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
                     del head_idx
-                    q_id = sequence_id[batch_idx, q_idx]
-                    kv_id = sequence_id[batch_idx, kv_idx]
+                    q_id = mask_pattern[batch_idx, q_idx]
+                    kv_id = mask_pattern[batch_idx, kv_idx]
                     return q_id == kv_id
 
             flex_block_mask = _get_flex_block_mask(
-                mask_pattern=sequence_id,
+                mask_pattern=mask_pattern,
                 batch_size=batch_size,
                 query_length=seq_len,
                 key_value_length=seq_len,
@@ -859,7 +1015,7 @@ class TransformerStack(nn.Module):
                 dtype=dtype,
                 mask_semantics=(
                     "boolean_sequence_id"
-                    if sequence_id.dtype == torch.bool
+                    if mask_pattern.dtype == torch.bool
                     else "integer_sequence_id"
                 ),
                 mask_mod=mask_mod,
@@ -889,6 +1045,237 @@ class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel):
         "flash_attention_3",
     )
 
+    def __init__(self, config: ESMplusplusConfig, *args: object, **kwargs: object) -> None:
+        super().__init__(config, *args, **kwargs)
+        self._sae_models = nn.ModuleDict()
+        self._esmc_fp8 = False
+        self._esmc_fp8_module_paths: tuple[str, ...] = ()
+        self._esmc_precision_status = ESMplusplusFP8Status(
+            enabled=False,
+            reason="FP8 has not been enabled; canonical checkpoint precision is unchanged.",
+            device="cpu",
+            transformer_engine_version=_transformer_engine_version(),
+            converted_projections=0,
+        )
+
+    @property
+    def esmc_precision_status(self) -> ESMplusplusFP8Status:
+        """Return the explicit FP8 conversion status for this ESM++ instance."""
+
+        return self._esmc_precision_status
+
+    def enable_fp8(self) -> ESMplusplusFP8Status:
+        """Enable the strict inference-only Transformer Engine FP8 path.
+
+        Canonical parameters must already be BF16 on one supported CUDA device.
+        Exactly one attention output projection per transformer block is replaced;
+        all other operations and all SAE weights remain BF16.
+        """
+
+        if self._esmc_fp8:
+            return self._esmc_precision_status
+        if self.training:
+            raise RuntimeError("ESM++ FP8 is inference-only; call eval() before enable_fp8().")
+        parameter_devices = {parameter.device for parameter in self.parameters()}
+        if len(parameter_devices) != 1:
+            raise RuntimeError(
+                "ESM++ FP8 requires every parameter on one CUDA device; sharded device maps "
+                f"are unsupported, found {sorted(map(str, parameter_devices))}."
+            )
+        device = next(iter(parameter_devices), self.device)
+        available, reason = _te_fp8_capability(device)
+        if not available:
+            raise RuntimeError(f"ESM++ FP8 is unavailable: {reason}")
+        non_bf16 = [
+            name
+            for name, parameter in self.named_parameters()
+            if parameter.is_floating_point() and parameter.dtype != torch.bfloat16
+        ]
+        if non_bf16:
+            examples = ", ".join(non_bf16[:3])
+            raise RuntimeError(
+                "ESM++ FP8 requires canonical BF16 parameters before conversion; "
+                f"found {len(non_bf16)} non-BF16 parameters (for example: {examples})."
+            )
+        paths = _convert_esmc_attention_outputs_to_te(
+            self,
+            expected_projections=self.config.num_hidden_layers,
+        )
+        self._esmc_fp8 = True
+        self._esmc_fp8_module_paths = paths
+        self._esmc_precision_status = ESMplusplusFP8Status(
+            enabled=True,
+            reason=(
+                f"{reason} Converted {len(paths)} attention output projections; "
+                "canonical checkpoint and SAE weights remain BF16."
+            ),
+            device=str(device),
+            transformer_engine_version=_transformer_engine_version(),
+            converted_projections=len(paths),
+        )
+        return self._esmc_precision_status
+
+    def add_sae_models(self, sae_models: list[nn.Module]) -> None:
+        """Attach official Biohub hidden-state SAE layers to this ESM++ model."""
+
+        for sae_model in sae_models:
+            if not isinstance(sae_model, nn.Module):
+                raise TypeError(
+                    "Each SAE must be an nn.Module obtained from an official Biohub "
+                    "ESMCSAEModel.layers entry."
+                )
+            layer = getattr(sae_model, "layer", None)
+            if isinstance(layer, bool) or not isinstance(layer, int):
+                raise TypeError("Each SAE layer must expose an integer .layer attribute.")
+            if not 0 <= layer <= self.config.num_hidden_layers:
+                raise ValueError(
+                    f"SAE target layer {layer} is outside the ESM++ hidden-state range "
+                    f"0..{self.config.num_hidden_layers}."
+                )
+            params = getattr(sae_model, "params", None)
+            d_model = getattr(params, "d_model", None)
+            if d_model != self.config.hidden_size:
+                raise ValueError(
+                    f"SAE layer {layer} expects d_model={d_model!r}, but this ESM++ "
+                    f"checkpoint has hidden_size={self.config.hidden_size}."
+                )
+            if not callable(getattr(sae_model, "get_sae_output", None)):
+                raise TypeError("Each SAE layer must expose get_sae_output(layer_states, token_mask).")
+            for name in ("idf", "max"):
+                if not isinstance(getattr(sae_model, name, None), torch.Tensor):
+                    raise TypeError(f"Each SAE layer must expose a tensor {name!r} buffer.")
+            key = f"layer{layer}"
+            if key in self._sae_models:
+                raise ValueError(
+                    f"An SAE is already registered at {key!r}; only one SAE per layer "
+                    "can be active."
+                )
+            self._sae_models[key] = sae_model
+
+    def _prepare_sae_forward(
+        self,
+        *,
+        compute_sae: bool,
+        input_ids: torch.Tensor | None,
+        attention_mask: torch.Tensor | None,
+        sequence_id: torch.Tensor | None,
+    ) -> tuple[tuple[int, ...], torch.Tensor | None]:
+        if not compute_sae or not self._sae_models:
+            return (), None
+        if input_ids is None:
+            raise ValueError(
+                "SAE computation requires input_ids so masked-token inputs can be rejected."
+            )
+        if torch.any(input_ids == self.config.mask_token_id):
+            raise ValueError("SAE inputs must not contain mask tokens; SAEs were trained unmasked.")
+        if sequence_id is not None:
+            token_mask = sequence_id >= 0
+        elif attention_mask is not None:
+            token_mask = attention_mask.to(dtype=torch.bool)
+        else:
+            token_mask = input_ids != self.config.pad_token_id
+        layers = tuple(sorted(int(name.removeprefix("layer")) for name in self._sae_models))
+        return layers, token_mask
+
+    def _get_sae_outputs(
+        self,
+        hidden_states: dict[int, torch.Tensor] | None,
+        token_mask: torch.Tensor | None,
+        *,
+        normalize_sae: bool,
+    ) -> dict[str, torch.Tensor] | None:
+        if not self._sae_models:
+            return None
+        if hidden_states is None or token_mask is None:
+            raise RuntimeError("SAE hidden-state collection was not initialized.")
+        outputs: dict[str, torch.Tensor] = {}
+        for key, sae_model in self._sae_models.items():
+            layer = int(key.removeprefix("layer"))
+            if layer not in hidden_states:
+                raise RuntimeError(f"ESM++ did not collect the requested SAE layer {layer}.")
+            sae_output = sae_model.get_sae_output(hidden_states[layer].clone(), token_mask)
+            features = getattr(sae_output, "feature_magnitudes", None)
+            if not isinstance(features, torch.Tensor):
+                raise TypeError("SAE get_sae_output must return tensor feature_magnitudes.")
+            features = features.detach()
+            if normalize_sae:
+                features = (features / sae_model.max) * sae_model.idf
+            outputs[key] = features.to_sparse()
+        return outputs
+
+    def _pad_fp8_inputs(
+        self,
+        input_ids: torch.Tensor | None,
+        attention_mask: torch.Tensor | None,
+        sequence_id: torch.Tensor | None,
+        inputs_embeds: torch.Tensor | None,
+    ) -> tuple[
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        int | None,
+    ]:
+        if not self._esmc_fp8:
+            return input_ids, attention_mask, sequence_id, inputs_embeds, None
+        source = input_ids if input_ids is not None else inputs_embeds
+        if source is None:
+            return input_ids, attention_mask, sequence_id, inputs_embeds, None
+        sequence_length = source.shape[1]
+        padded_length = (
+            (sequence_length + _ESMC_FP8_ALIGNMENT - 1) // _ESMC_FP8_ALIGNMENT
+        ) * _ESMC_FP8_ALIGNMENT
+        padding = padded_length - sequence_length
+        if padding == 0:
+            return input_ids, attention_mask, sequence_id, inputs_embeds, None
+        if input_ids is not None:
+            input_ids = F.pad(input_ids, (0, padding), value=self.config.pad_token_id)
+        if inputs_embeds is not None:
+            inputs_embeds = F.pad(inputs_embeds, (0, 0, 0, padding), value=0.0)
+        if sequence_id is not None:
+            sequence_id = F.pad(sequence_id.to(dtype=torch.long), (0, padding), value=-1)
+        else:
+            if attention_mask is None:
+                attention_mask = (
+                    source != self.config.pad_token_id
+                    if input_ids is not None
+                    else torch.ones(
+                        source.shape[:2],
+                        dtype=torch.bool,
+                        device=source.device,
+                    )
+                )
+            attention_mask = F.pad(attention_mask, (0, padding), value=0)
+        return input_ids, attention_mask, sequence_id, inputs_embeds, sequence_length
+
+    @staticmethod
+    def _trim_transformer_output(
+        output: TransformerOutput,
+        sequence_length: int | None,
+    ) -> TransformerOutput:
+        if sequence_length is None:
+            return output
+        hidden_states = (
+            tuple(state[:, :sequence_length] for state in output.hidden_states)
+            if output.hidden_states is not None
+            else None
+        )
+        attentions = (
+            tuple(
+                attention[..., :sequence_length, :sequence_length]
+                for attention in output.attentions
+            )
+            if output.attentions is not None
+            else None
+        )
+        return TransformerOutput(
+            last_hidden_state=output.last_hidden_state[:, :sequence_length],
+            hidden_states=hidden_states,
+            attentions=attentions,
+            s_max=output.s_max,
+            sae_hidden_states=output.sae_hidden_states,
+        )
+
     @property
     def tokenizer(self) -> EsmSequenceTokenizer:
         """Construct the sequence tokenizer only when a raw-sequence API needs it."""
@@ -1033,6 +1420,8 @@ class ESMplusplusModel(PreTrainedESMplusplusModel, EmbeddingMixin):
         output_s_max: bool | None = False,
         esmfold2_hidden_states: bool = False,
         return_dict: bool | None = None,
+        compute_sae: bool = True,
+        normalize_sae: bool = False,
     ) -> TransformerOutput | tuple[torch.Tensor, ...]:
         """Run ESMC inference with the pinned Biohub mask precedence.
 
@@ -1056,25 +1445,47 @@ class ESMplusplusModel(PreTrainedESMplusplusModel, EmbeddingMixin):
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
+        input_ids, attention_mask, sequence_id, inputs_embeds, original_length = (
+            self._pad_fp8_inputs(input_ids, attention_mask, sequence_id, inputs_embeds)
+        )
         if attention_mask is None and sequence_id is None and input_ids is not None:
             attention_mask = input_ids.ne(self.config.pad_token_id)
+        sae_layers, sae_token_mask = self._prepare_sae_forward(
+            compute_sae=compute_sae,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            sequence_id=sequence_id,
+        )
 
         x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds
 
-        transformer_output = self.transformer(
-            x=x,
-            attention_mask=attention_mask,
-            sequence_id=sequence_id,
-            output_hidden_states=output_hidden_states,
-            output_attentions=output_attentions,
-            output_s_max=output_s_max,
-            esmfold2_hidden_states=esmfold2_hidden_states,
+        with _esmplusplus_fp8_context(self._esmc_fp8, self.device):
+            transformer_output = self.transformer(
+                x=x,
+                attention_mask=attention_mask,
+                sequence_id=sequence_id,
+                output_hidden_states=output_hidden_states,
+                output_attentions=output_attentions,
+                output_s_max=output_s_max,
+                esmfold2_hidden_states=esmfold2_hidden_states,
+                sae_layers=sae_layers,
+            )
+        sae_outputs = (
+            self._get_sae_outputs(
+                transformer_output.sae_hidden_states,
+                sae_token_mask,
+                normalize_sae=normalize_sae,
+            )
+            if sae_layers
+            else None
         )
+        transformer_output = self._trim_transformer_output(transformer_output, original_length)
         result = TransformerOutput(
             last_hidden_state=transformer_output.last_hidden_state,
             hidden_states=transformer_output.hidden_states,
             attentions=transformer_output.attentions,
             s_max=transformer_output.s_max,
+            sae_outputs=sae_outputs,
         )
         return result if return_dict else result.to_tuple()
 
@@ -1158,6 +1569,8 @@ class ESMplusplusForMaskedLM(
         esmfold2_hidden_states: bool = False,
         return_dict: bool | None = None,
         compute_logits: bool = True,
+        compute_sae: bool = True,
+        normalize_sae: bool = False,
     ) -> ESMplusplusOutput | tuple[torch.Tensor, ...]:
         if input_ids is None and inputs_embeds is None:
             raise ValueError("You have to specify either input_ids or inputs_embeds")
@@ -1174,20 +1587,41 @@ class ESMplusplusForMaskedLM(
             else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        input_ids, attention_mask, sequence_id, inputs_embeds, original_length = (
+            self._pad_fp8_inputs(input_ids, attention_mask, sequence_id, inputs_embeds)
+        )
         if attention_mask is None and sequence_id is None and input_ids is not None:
             attention_mask = input_ids.ne(self.config.pad_token_id)
+        sae_layers, sae_token_mask = self._prepare_sae_forward(
+            compute_sae=compute_sae,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            sequence_id=sequence_id,
+        )
 
         x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds
 
-        output = self.transformer(
-            x=x,
-            attention_mask=attention_mask,
-            sequence_id=sequence_id,
-            output_hidden_states=output_hidden_states,
-            output_attentions=output_attentions,
-            output_s_max=output_s_max,
-            esmfold2_hidden_states=esmfold2_hidden_states,
+        with _esmplusplus_fp8_context(self._esmc_fp8, self.device):
+            output = self.transformer(
+                x=x,
+                attention_mask=attention_mask,
+                sequence_id=sequence_id,
+                output_hidden_states=output_hidden_states,
+                output_attentions=output_attentions,
+                output_s_max=output_s_max,
+                esmfold2_hidden_states=esmfold2_hidden_states,
+                sae_layers=sae_layers,
+            )
+        sae_outputs = (
+            self._get_sae_outputs(
+                output.sae_hidden_states,
+                sae_token_mask,
+                normalize_sae=normalize_sae,
+            )
+            if sae_layers
+            else None
         )
+        output = self._trim_transformer_output(output, original_length)
 
         last_hidden_state = output.last_hidden_state
         logits = self.sequence_head(last_hidden_state) if compute_logits else None
@@ -1205,6 +1639,7 @@ class ESMplusplusForMaskedLM(
             attentions=output.attentions,
             s_max=output.s_max,
             last_hidden_state=last_hidden_state,
+            sae_outputs=sae_outputs,
         )
         return result if return_dict else result.to_tuple()
 
@@ -1280,6 +1715,8 @@ class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM, EmbeddingMixi
         output_hidden_states: bool | None = None,
         output_s_max: bool | None = False,
         return_dict: bool | None = None,
+        compute_sae: bool = True,
+        normalize_sae: bool = False,
     ) -> ESMplusplusSequenceClassifierOutput | tuple[torch.Tensor, ...]:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         pooling_mask = attention_mask
@@ -1310,6 +1747,8 @@ class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM, EmbeddingMixi
             output_s_max=output_s_max,
             return_dict=True,
             compute_logits=False,
+            compute_sae=compute_sae,
+            normalize_sae=normalize_sae,
         )
 
         last_hidden_state = output.last_hidden_state
@@ -1345,6 +1784,7 @@ class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM, EmbeddingMixi
             hidden_states=output.hidden_states,
             attentions=output.attentions,
             s_max=output.s_max,
+            sae_outputs=output.sae_outputs,
         )
         return result if return_dict else result.to_tuple()
 
@@ -1399,6 +1839,8 @@ class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM, EmbeddingMixin):
         output_hidden_states: bool | None = None,
         output_s_max: bool | None = False,
         return_dict: bool | None = None,
+        compute_sae: bool = True,
+        normalize_sae: bool = False,
     ) -> ESMplusplusTokenClassifierOutput | tuple[torch.Tensor, ...]:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         output = super().forward(
@@ -1412,6 +1854,8 @@ class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM, EmbeddingMixin):
             output_s_max=output_s_max,
             return_dict=True,
             compute_logits=False,
+            compute_sae=compute_sae,
+            normalize_sae=normalize_sae,
         )
 
         last_hidden_state = output.last_hidden_state
@@ -1427,6 +1871,7 @@ class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM, EmbeddingMixin):
             hidden_states=output.hidden_states,
             attentions=output.attentions,
             s_max=output.s_max,
+            sae_outputs=output.sae_outputs,
         )
         return result if return_dict else result.to_tuple()
 

fastplms/models/esmfold2/esmfold2_processor.py

@@ -123,12 +123,6 @@ def clean_esmfold2_input(input: StructurePredictionInput) -> StructurePrediction
             "the published ESMFold2 feature pipeline drops it. FastPLMs refuses this "
             "input instead of silently emitting an all-zero pocket feature."
         )
-    if input.distogram_conditioning is not None:
-        raise NotImplementedError(
-            "ESMFold2 distogram conditioning is present in the upstream input schema but "
-            "the published ESMFold2 forward does not consume it. FastPLMs refuses this "
-            "input instead of silently ignoring the supplied distogram."
-        )
 
     cleaned: list[Any] = []
     for item in input.sequences:

fastplms/models/esmfold2/modeling_esmfold2.py

@@ -74,7 +74,6 @@ from .modeling_esmfold2_common import (
     maybe_subsample_msa,
     validate_kernel_backend,
     validate_msa_conditioning_inputs,
-    validate_prepared_auxiliary_inputs,
 )
 
 _ESMC_FP8_LINEAR_SUFFIX = ".attn.out_proj"
@@ -1418,12 +1417,6 @@ class ESMFold2Model(
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
-        pocket_feature: Tensor | None = None,
-        gt_coords: Tensor | None = None,
-        is_resolved: Tensor | None = None,
-        frames_idx: Tensor | None = None,
-        disto_cond: Tensor | None = None,
-        disto_cond_mask: Tensor | None = None,
     ) -> ESMFold2Output | tuple[Any, ...]:
         output_hidden_states, return_dict = _resolve_structure_output_controls(
             self.config,
@@ -1439,12 +1432,6 @@ class ESMFold2Model(
             deletion_value=deletion_value,
             deletion_mean=deletion_mean,
         )
-        validate_prepared_auxiliary_inputs(
-            pocket_feature=pocket_feature,
-            disto_cond=disto_cond,
-            disto_cond_mask=disto_cond_mask,
-        )
-        del gt_coords, is_resolved, frames_idx
         tok_mask = token_attention_mask
         atm_mask = atom_attention_mask
         disto_idx = distogram_atom_idx

fastplms/models/esmfold2/modeling_esmfold2_common.py

@@ -57,14 +57,6 @@ MSA_CONDITIONING_INPUT_NAMES = (
     "deletion_value",
     "deletion_mean",
 )
-PREPARED_AUXILIARY_INPUT_NAMES = (
-    "pocket_feature",
-    "gt_coords",
-    "is_resolved",
-    "frames_idx",
-    "disto_cond",
-    "disto_cond_mask",
-)
 
 
 def validate_kernel_backend(backend: str | None) -> None:
@@ -113,29 +105,6 @@ def validate_msa_conditioning_inputs(
         )
 
 
-def validate_prepared_auxiliary_inputs(
-    *,
-    pocket_feature: Tensor | None,
-    disto_cond: Tensor | None,
-    disto_cond_mask: Tensor | None,
-) -> None:
-    """Accept inert upstream features while rejecting unsupported conditioning."""
-
-    if pocket_feature is not None and torch.any(pocket_feature != 0).item():
-        raise NotImplementedError(
-            "The published ESMFold2 forward does not consume pocket conditioning; "
-            "nonzero pocket_feature values are unsupported."
-        )
-    distogram_is_active = (disto_cond is not None and torch.any(disto_cond != 0).item()) or (
-        disto_cond_mask is not None and torch.any(disto_cond_mask).item()
-    )
-    if distogram_is_active:
-        raise NotImplementedError(
-            "The published ESMFold2 forward does not consume distogram conditioning; "
-            "nonzero disto_cond or disto_cond_mask values are unsupported."
-        )
-
-
 def _fused_active(module: nn.Module, tensor: Tensor) -> bool:
     """Return whether an optional fused implementation can handle this call."""
     return (
@@ -1612,6 +1581,22 @@ class SWAAtomTransformer(nn.Module):
 # ===========================================================================
 
 
+@torch.compiler.disable
+def _prepare_atom_encoder_metadata(
+    atom_attention_mask: Tensor,
+    atom_to_token: Tensor,
+    num_diffusion_samples: int,
+) -> tuple[Tensor, Tensor, Tensor, int, int]:
+    """Prepare mask-derived atom metadata outside compiled diffusion graphs."""
+    mask_exp = atom_attention_mask.repeat_interleave(num_diffusion_samples, 0)
+    seqlens = mask_exp.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(mask_exp.flatten(), as_tuple=False).flatten()
+    max_seqlen = int(seqlens.max().item())
+    cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
+    n_tokens = int(atom_to_token.max().item()) + 1
+    return mask_exp, indices, cu_seqlens, max_seqlen, n_tokens
+
+
 class ESMFold2AtomEncoder(nn.Module):
     """Encode atom inputs with normalization and sliding-window attention.
 
@@ -1708,13 +1693,14 @@ class ESMFold2AtomEncoder(nn.Module):
             cos, sin = self.atom_transformer._build_3d_rope(ref_pos, ref_space_uid)
             cos = cos.repeat_interleave(num_diffusion_samples, 0)
             sin = sin.repeat_interleave(num_diffusion_samples, 0)
-            mask_exp = atom_attention_mask.repeat_interleave(num_diffusion_samples, 0)
-            seqlens = mask_exp.sum(dim=-1, dtype=torch.int32)
-            indices = torch.nonzero(mask_exp.flatten(), as_tuple=False).flatten()
-            max_seqlen = int(seqlens.max().item())
-            cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
+            mask_exp, indices, cu_seqlens, max_seqlen, n_tokens = (
+                _prepare_atom_encoder_metadata(
+                    atom_attention_mask,
+                    atom_to_token,
+                    num_diffusion_samples,
+                )
+            )
             attention_params = (cos, sin, indices, cu_seqlens, max_seqlen)
-            n_tokens = int(atom_to_token.max().item()) + 1
             if layer_cache is not None:
                 layer_cache["c_base"] = c_base
                 layer_cache["attention_params"] = attention_params

fastplms/models/esmfold2/modeling_esmfold2_experimental.py

@@ -59,7 +59,6 @@ from .modeling_esmfold2_common import (
     gather_token_to_atom,
     validate_kernel_backend,
     validate_msa_conditioning_inputs,
-    validate_prepared_auxiliary_inputs,
 )
 
 _EPS = 1e-5
@@ -690,12 +689,6 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
-        pocket_feature: Tensor | None = None,
-        gt_coords: Tensor | None = None,
-        is_resolved: Tensor | None = None,
-        frames_idx: Tensor | None = None,
-        disto_cond: Tensor | None = None,
-        disto_cond_mask: Tensor | None = None,
     ) -> ESMFold2Output | tuple[Any, ...]:
         output_hidden_states, return_dict = _resolve_structure_output_controls(
             self.config,
@@ -711,12 +704,6 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             deletion_value=deletion_value,
             deletion_mean=deletion_mean,
         )
-        validate_prepared_auxiliary_inputs(
-            pocket_feature=pocket_feature,
-            disto_cond=disto_cond,
-            disto_cond_mask=disto_cond_mask,
-        )
-        del gt_coords, is_resolved, frames_idx
         tok_mask = token_attention_mask
         atm_mask = atom_attention_mask
         n_loops = num_loops if num_loops is not None else self.config.num_loops

modeling_fastplms.py

@@ -12,7 +12,7 @@ from zipfile import ZIP_DEFLATED, ZipFile
 
 from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
 
-if RUNTIME_HASH != "6e81dad15aa6df4f53633997daf8747aa49a4afbd030a65937a6f3a7d8159273":
+if RUNTIME_HASH != "5f557b57da6c13e8b26b2a144ab0983b9621d853a3141a3a52e912d2de26289a":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []