Update FastPLMs files

README.md

@@ -10,7 +10,7 @@ tags:
 
 # Synthyra/ESMplusplus_large
 
-This checkpoint packages the FastPLMs `ESMC` implementation.
+This checkpoint contains the FastPLMs `ESMC` implementation.
 
 Accepted inputs are amino-acid sequences tokenized to residue IDs.
 Supported Transformers entry points are `AutoConfig`, `AutoModel`,
@@ -28,9 +28,7 @@ Supported Transformers entry points are `AutoConfig`, `AutoModel`,
 | Attention variants | Special: SDPA fidelity path; alternate backends have explicit bands |
 | 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/ESMplusplus_large/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 FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. 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 FlashAttention loader dependency. FlashAttention also requires compatible CUDA hardware and BF16 execution. 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,23 +60,23 @@ model = AutoModel.from_pretrained(
 ```
 
 For offline validation, replace `model_id` with the manifest-built
-`dist/hub/ESMplusplus_large` path and pass `local_files_only=True`.
+`dist/hub/ESMplusplus_large` path. Pass `local_files_only=True`.
 
 ## Attention and compliance
 
 The quick start selects `sdpa` explicitly. Declared variants are `eager`, `sdpa`, `flex_attention`, `flash_attention_2`,
-`flash_attention_3`. 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.
+`flash_attention_3`. An unavailable 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.
 
 ## Tokenization and forward inference
 
-Load the tokenizer from the same artifact as the model. Padding is represented
-explicitly by the attention mask:
+Load the tokenizer from the same artifact as the model. The attention mask
+shows padding explicitly:
 
 ```python
 import torch
@@ -103,8 +101,8 @@ print(output.last_hidden_state.shape)
 
 ## Dataset embeddings
 
-The shared embedding mixin preserves input order and biological-position
-masking. It accepts sequences, identified records, mappings, or a FASTA path:
+The shared embedding mixin keeps input order and biological-position masking.
+It accepts sequences, identified records, mappings, or a FASTA path:
 
 ```python
 pooled = model.embed_dataset(
@@ -121,12 +119,12 @@ print(residues[0].tensor.shape) # (l, d)
 ```
 
 Set `output` and `format="safetensors"` or `"sqlite"` for transactional,
-bounded-memory persistence. Resume verifies input order, model state, tokenizer
-policy, backend, dtype, and pooling configuration before appending.
+bounded-memory storage. Resume checks input order, model state, tokenizer
+policy, backend, dtype, and pooling configuration before it appends data.
 
 ## 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"
@@ -145,17 +143,17 @@ 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.
 
 ## Test-time training
 
 TTT samples masked views of one protein and updates only injected low-rank
-adapters. Base checkpoint weights remain frozen:
+adapters. Base checkpoint weights stay frozen:
 
 ```python
 from transformers import AutoModelForMaskedLM
@@ -173,28 +171,85 @@ ttt_model.ttt_reset()
 print(metrics)
 ```
 
-Persisted adapters retain their deterministic reset state. TTT adds latency
-and memory, can worsen an output, and does not establish biological function.
+Saved adapters retain their deterministic reset state. TTT adds latency and
+memory, can worsen an output, and does not show biological function.
 
 ## ESMC behavior
 
-This artifact exposes the Biohub ESMC sequence encoder and masked-language-model
-head through Transformers. It is also the language-model family used by
-ESMFold2. SDPA is the default and the recommended choice for highest numerical
-fidelity. Flex Attention and FlashAttention 3 are supported, non-experimental
-backends, but their BF16 arithmetic may be numerically divergent from SDPA.
-Those deviations produce diagnostic warnings rather than strict parity
-failures; dispatch integrity, masks, finite outputs, shapes, and catastrophic
-biological disagreement remain hard gates.
+This artifact provides the Biohub ESMC sequence encoder and masked-language-
+model head through Transformers. ESMFold2 also uses this language-model family.
+SDPA is the default and gives the highest numerical fidelity. Flex Attention and
+FlashAttention 3 are supported non-experimental backends. Their BF16 arithmetic
+can differ numerically from SDPA. These differences give diagnostic warnings,
+not strict-parity failures. Dispatch, masks, finite outputs, shapes, and large
+biological disagreements remain hard gates.
 
 The current GH200/aarch64 release environment validates eager, SDPA, and Flex.
-Flash requests fail closed because compatible locked kernels are unavailable
-on this platform.
+Flash requests raise because compatible locked kernels are unavailable on this
+platform.
 
-When `sequence_id` is supplied, it is authoritative for ESMC attention grouping
-and padding, and `attention_mask` is ignored. Values greater than or equal to
-zero are valid sequence-group IDs; `-1` denotes padding. Omit `sequence_id` to
-use `attention_mask` as the padding contract.
+When `sequence_id` is supplied, it controls ESMC attention groups and padding.
+`attention_mask` is ignored. Values greater than or equal to zero are valid
+sequence-group IDs. `-1` marks padding. Omit `sequence_id` to use
+`attention_mask` for padding.
+
+### Hidden-state sparse autoencoders
+
+ESM++ supports hidden-state SAEs from the official
+[Biohub ESMC SAE collection](https://huggingface.co/collections/biohub/esmc-saes-for-hidden-states-all-layers).
+Select an SAE for this ESMC scale. Load only required layers. Then attach them
+to the model:
+
+```python
+import torch
+from transformers import AutoModel
+
+sae = AutoModel.from_pretrained("biohub/ESMC-600M-sae-layer27-k64-codebook65536", device=model.device)
+sae.initialize_layers([27])
+model.add_sae_models([sae.layers["27"]])
+
+with torch.inference_mode():
+    output = model(**batch, normalize_sae=True)
+
+features = output.sae_outputs["layer27"]
+print(features.shape, features.layout)  # (valid_token_count, codebook_dim), sparse COO
+```
+
+SAEs run after you attach them. Use `compute_sae=False` to skip SAE work.
+Outputs are detached sparse tensors with keys such as `layer{N}`. They omit
+padding. The model uses `sequence_id`, then `attention_mask`, for padding.
+`normalize_sae=True` uses Biohub `(features / max) * idf` normalization. SAE
+computation requires `input_ids`. It rejects mask tokens because Biohub trained
+the SAEs with unmasked sequences. This interface supports hidden-state SAEs
+only, not MLP-output SAEs. FastPLMs does not copy SAE weights or add SAE
+checkpoints to its model manifest.
+
+### Experimental FP8 inference
+
+The default uses checkpoint BF16 behavior. FP8 is an explicit experimental
+inference option for every ESM++ scale:
+
+```python
+import torch
+from transformers import AutoModel
+
+fp8_model = AutoModel.from_pretrained(
+    "Synthyra/ESMplusplus_large",
+    trust_remote_code=True,
+    dtype=torch.bfloat16,
+).cuda().eval()
+fp8_model.enable_fp8()
+print(fp8_model.esmc_precision_status)
+
+with torch.inference_mode():
+    fp8_output = fp8_model(**{name: value.cuda() for name, value in batch.items()})
+```
+
+FP8 forward calls require `torch.inference_mode()`. The model pads the sequence
+dimension to a multiple of 16. Transformer Engine converts supported linear
+layers. The call fails if the dependency, compatible CUDA hardware, or complete
+conversion set is unavailable. It does not silently use BF16. FP8 does not
+claim numerical parity.
 
 | Backend | Support | Measurement status |
 | --- | --- | --- |
@@ -217,7 +272,7 @@ and
 - Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`
 - AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`
 - Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3`
-- Precision policies: `default`
+- Precision policies: `default`, `fp8` (experimental)
 - BF16 execution: `static_parameters`
 - Generation contract: `not_applicable`
 - Artifact dependency set: `core`
@@ -229,8 +284,8 @@ and
 ## Release record
 
 - FastPLMs weights: `Synthyra/ESMplusplus_large`
-- 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/ESMC-600M`
 - Artifact source: `fast`
 - State transform: `esmc_to_fastplms_v1`
@@ -238,19 +293,17 @@ and
 - Release tiers: `check`, `compliance`, `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()
 

modeling_fastplms.py

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