Update FastPLMs files

README.md

@@ -10,7 +10,7 @@ tags:
 
 # Synthyra/ESM3_small
 
-This checkpoint packages the FastPLMs `ESM3` implementation.
+This checkpoint contains the FastPLMs `ESM3` implementation.
 
 Accepted inputs are sequence, structure, and function tracks prepared through
 the multimodal helpers.
@@ -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/ESM3_small/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 declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. 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 CPU gate covers small offline tests. Published checkpoint throughput and parity require the documented device tier. 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,22 +60,22 @@ model = AutoModel.from_pretrained(
 ```
 
 For offline validation, replace `model_id` with the manifest-built
-`dist/hub/ESM3_small` path and pass `local_files_only=True`.
+`dist/hub/ESM3_small` 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.
 
 ## 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(
@@ -94,12 +92,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"
@@ -118,17 +116,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 AutoModel
@@ -146,14 +144,14 @@ 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.
 
 ## Sequence inference and masked-sequence generation
 
-ESM3 owns its sequence preparation. This example exercises the sequence track;
-the public input contract also supports structure and function tracks through
-the multimodal helpers:
+ESM3 prepares its sequence input. This example uses the sequence track. The
+public input contract also supports structure and function tracks through the
+multimodal helpers:
 
 ```python
 import torch
@@ -171,10 +169,9 @@ print(output.structure_logits.shape)
 print(output.function_logits.shape)
 ```
 
-When `return_dict=False`, ESM3 follows the standard base-model tuple prefix:
+When `return_dict=False`, ESM3 uses the standard base-model tuple prefix:
 `last_hidden_state`, then requested `hidden_states` and `attentions`. Multimodal
-logits and extensions follow that prefix. Prefer named fields for individual
-tracks.
+logits and extensions follow this prefix. Use named fields for individual tracks.
 
 Generate masked sequence positions with an explicit seed:
 
@@ -190,8 +187,8 @@ generated = model.generate("MK____A", config)
 print(generated)
 ```
 
-Underscores mark positions to generate. Model outputs are predictions over
-tracks, not experimental measurements of structure or function.
+Underscores mark positions to generate. Model outputs are track predictions,
+not experimental measurements of structure or function.
 
 ## Runtime contract
 
@@ -211,8 +208,8 @@ tracks, not experimental measurements of structure or function.
 ## Release record
 
 - FastPLMs weights: `Synthyra/ESM3_small`
-- 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/esm3-sm-open-v1`
 - Artifact source: `fast`
 - State transform: `esm3_to_fastplms_v1`
@@ -220,19 +217,17 @@ tracks, not experimental measurements of structure or function.
 - 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/esm3/modeling_esm3.py

@@ -1350,12 +1350,13 @@ class MultiHeadAttention(nn.Module):
     def forward(
         self,
         x: torch.Tensor,
-        seq_id: torch.Tensor | None,
         attention_mask: torch.Tensor | None = None,
+        flex_block_mask: BlockMask | None = None,
+        mask_semantics: str = "dense",
         output_attentions: bool = False,
         effective_backend: AttentionBackend | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        # x: (b, l, d); seq_id, attention_mask: (b, l)
+        # x: (b, l, d); attention_mask: (b, 1, l, l) or (b, 1, 1, l)
         qkv = self.layernorm_qkv(x)  # (b, l, 3 * d)
         query, key, value = torch.chunk(qkv, 3, dim=-1)
         query = self.q_ln(query).to(query.dtype)
@@ -1368,13 +1369,7 @@ class MultiHeadAttention(nn.Module):
             h=self.n_heads,
         )
         query, key, value = map(reshaper, (query, key, value))  # each (b, h, l, d_h)
-
-        mask = None
-        if seq_id is not None:
-            mask = (seq_id.unsqueeze(-1) == seq_id.unsqueeze(-2)).unsqueeze(1)
-        if attention_mask is not None:
-            key_padding_mask = attention_mask[:, None, None, :]
-            mask = key_padding_mask if mask is None else mask & key_padding_mask
+        mask = attention_mask
 
         if effective_backend is None:
             effective_backend = resolve_attention_backend_for_call(
@@ -1401,15 +1396,6 @@ class MultiHeadAttention(nn.Module):
         else:
             attn_weights = None
             if effective_backend == AttentionBackend.FLEX:
-                block_mask = self._create_flex_block_mask(seq_id, attention_mask, query)
-                if seq_id is not None and attention_mask is not None:
-                    mask_semantics = "sequence_id_and_padding"
-                elif seq_id is not None:
-                    mask_semantics = "sequence_id_equality"
-                elif attention_mask is not None:
-                    mask_semantics = "padding"
-                else:
-                    mask_semantics = "dense"
                 fn = _get_flex_attention_fn(
                     device=query.device,
                     dtype=query.dtype,
@@ -1423,7 +1409,7 @@ class MultiHeadAttention(nn.Module):
                     query,
                     key,
                     value,
-                    block_mask=block_mask,
+                    block_mask=flex_block_mask,
                     scale=self.scale,
                 )
             elif effective_backend == AttentionBackend.SDPA:
@@ -1442,37 +1428,6 @@ class MultiHeadAttention(nn.Module):
         context = einops.rearrange(context, "b h s d -> b s (h d)")  # (b, l, d)
         return self.out_proj(context), attn_weights
 
-    @staticmethod
-    def _create_flex_block_mask(
-        seq_id: torch.Tensor | None,
-        attention_mask: torch.Tensor | None,
-        query: torch.Tensor,
-    ) -> BlockMask | None:
-        if seq_id is None and attention_mask is None:
-            return None
-        if create_block_mask is None:
-            raise RuntimeError(
-                "Flex Attention requested but torch.create_block_mask is unavailable."
-            )
-        batch_size, _, seq_len, _ = query.shape
-
-        def mask_mod(batch_idx, _head_idx, q_idx, kv_idx):
-            if seq_id is None:
-                return attention_mask[batch_idx, kv_idx]
-            allowed = seq_id[batch_idx, q_idx] == seq_id[batch_idx, kv_idx]
-            if attention_mask is not None:
-                allowed = allowed & attention_mask[batch_idx, kv_idx]
-            return allowed
-
-        return create_block_mask(
-            mask_mod,
-            batch_size,
-            1,
-            seq_len,
-            seq_len,
-            device=query.device,
-        )
-
 
 class GeometricReasoningOriginalImpl(nn.Module):
     def __init__(
@@ -1683,6 +1638,8 @@ class UnifiedTransformerBlock(nn.Module):
         x: torch.Tensor,
         sequence_id: torch.Tensor | None,
         attention_mask: torch.Tensor | None,
+        flex_block_mask: BlockMask | None,
+        mask_semantics: str,
         frames: Affine3D,
         frames_mask: torch.Tensor,
         chain_id: torch.Tensor,
@@ -1693,8 +1650,9 @@ class UnifiedTransformerBlock(nn.Module):
         if self.use_plain_attn:
             plain_residual, attn_weights = self.attn(
                 x,
-                sequence_id,
                 attention_mask,
+                flex_block_mask,
+                mask_semantics,
                 output_attentions=output_attentions,
                 effective_backend=effective_backend,
             )
@@ -1772,9 +1730,17 @@ class TransformerStack(nn.Module):
             chain_id = torch.ones(size=batch_dims, dtype=torch.int64, device=x.device)
         if affine is None or affine_mask is None:
             raise ValueError("affine and affine_mask are required for ESM3 transformer calls.")
-        effective_backend = resolve_attention_backend_for_call(
-            self.attention_backend,
-            output_attentions=output_attentions,
+        attention_mask, flex_block_mask, affine_mask, mask_semantics, effective_backend = (
+            self._prepare_attention_masks(
+                sequence_id=sequence_id,
+                attention_mask=attention_mask,
+                affine_mask=affine_mask,
+                batch_size=x.shape[0],
+                seq_len=x.shape[1],
+                device=x.device,
+                attention_backend=self.attention_backend,
+                output_attentions=output_attentions,
+            )
         )
         all_hidden_states = [] if output_hidden_states else None
         all_attentions = []
@@ -1783,6 +1749,8 @@ class TransformerStack(nn.Module):
                 x,
                 sequence_id,
                 attention_mask,
+                flex_block_mask,
+                mask_semantics,
                 affine,
                 affine_mask,
                 chain_id,
@@ -1797,6 +1765,110 @@ class TransformerStack(nn.Module):
         attentions = tuple(all_attentions) if output_attentions else None
         return self.norm(x), x, hidden_states, attentions
 
+    @staticmethod
+    @torch.compiler.disable
+    def _prepare_attention_masks(
+        sequence_id: torch.Tensor | None,
+        attention_mask: torch.Tensor | None,
+        affine_mask: torch.Tensor,
+        batch_size: int,
+        seq_len: int,
+        device: torch.device,
+        attention_backend: AttentionBackend,
+        output_attentions: bool,
+    ) -> tuple[
+        torch.Tensor | None,
+        BlockMask | None,
+        torch.Tensor,
+        str,
+        AttentionBackend,
+    ]:
+        expected_mask_shape = (batch_size, seq_len)
+        if sequence_id is not None and tuple(sequence_id.shape) != expected_mask_shape:
+            raise ValueError(
+                "sequence_id must have shape (batch, sequence); "
+                f"expected {expected_mask_shape}, received {tuple(sequence_id.shape)}."
+            )
+        if attention_mask is not None:
+            if tuple(attention_mask.shape) != expected_mask_shape:
+                raise ValueError(
+                    "attention_mask must have shape (batch, sequence); "
+                    f"expected {expected_mask_shape}, received {tuple(attention_mask.shape)}."
+                )
+            if attention_mask.dtype != torch.bool and not bool(
+                torch.logical_or(attention_mask == 0, attention_mask == 1).all()
+            ):
+                raise ValueError("attention_mask must contain only boolean or 0/1 values.")
+            attention_mask = attention_mask.to(device=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.")
+            affine_mask = affine_mask & attention_mask
+
+        effective_backend = resolve_attention_backend_for_call(
+            attention_backend,
+            output_attentions=output_attentions,
+        )
+
+        if sequence_id is not None and attention_mask is not None:
+            mask_semantics = "sequence_id_and_padding"
+        elif sequence_id is not None:
+            mask_semantics = "sequence_id_equality"
+        elif attention_mask is not None:
+            mask_semantics = "padding"
+        else:
+            mask_semantics = "dense"
+
+        dense_mask = None
+        flex_block_mask = None
+        has_attention_mask = sequence_id is not None or attention_mask is not None
+        if effective_backend == AttentionBackend.FLEX and has_attention_mask:
+            flex_block_mask = TransformerStack._create_flex_block_mask(
+                sequence_id,
+                attention_mask,
+                batch_size,
+                seq_len,
+                device,
+            )
+        else:
+            if sequence_id is not None:
+                dense_mask = (sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2)).unsqueeze(1)
+            if attention_mask is not None:
+                key_padding_mask = attention_mask[:, None, None, :]
+                dense_mask = (
+                    key_padding_mask if dense_mask is None else dense_mask & key_padding_mask
+                )
+        return dense_mask, flex_block_mask, affine_mask, mask_semantics, effective_backend
+
+    @staticmethod
+    def _create_flex_block_mask(
+        sequence_id: torch.Tensor | None,
+        attention_mask: torch.Tensor | None,
+        batch_size: int,
+        seq_len: int,
+        device: torch.device,
+    ) -> BlockMask:
+        if create_block_mask is None:
+            raise RuntimeError(
+                "Flex Attention requested but torch.create_block_mask is unavailable."
+            )
+
+        def mask_mod(batch_idx, _head_idx, q_idx, kv_idx):
+            if sequence_id is None:
+                return attention_mask[batch_idx, kv_idx]
+            allowed = sequence_id[batch_idx, q_idx] == sequence_id[batch_idx, kv_idx]
+            if attention_mask is not None:
+                allowed = allowed & attention_mask[batch_idx, kv_idx]
+            return allowed
+
+        return create_block_mask(
+            mask_mod,
+            batch_size,
+            1,
+            seq_len,
+            seq_len,
+            device=device,
+        )
+
 
 class EncodeInputs(nn.Module):
     def __init__(self, d_model: int, sequence_vocab_size: int = 64) -> None:
@@ -2049,26 +2121,6 @@ class ESM3Core(nn.Module):
             function_tokens,
             residue_annotation_tokens,
         )
-        expected_mask_shape = tuple(x.shape[:2])
-        if sequence_id is not None and tuple(sequence_id.shape) != expected_mask_shape:
-            raise ValueError(
-                "sequence_id must have shape (batch, sequence); "
-                f"expected {expected_mask_shape}, received {tuple(sequence_id.shape)}."
-            )
-        if attention_mask is not None:
-            if tuple(attention_mask.shape) != expected_mask_shape:
-                raise ValueError(
-                    "attention_mask must have shape (batch, sequence); "
-                    f"expected {expected_mask_shape}, received {tuple(attention_mask.shape)}."
-                )
-            if attention_mask.dtype != torch.bool and not bool(
-                torch.logical_or(attention_mask == 0, attention_mask == 1).all()
-            ):
-                raise ValueError("attention_mask must contain only boolean or 0/1 values.")
-            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.")
-            affine_mask = affine_mask & attention_mask
         x, embedding, hidden_states, attentions = self.transformer(
             x,
             sequence_id,

modeling_fastplms.py

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