Sync OSS esm 3.4.2

README.md

@@ -55,31 +55,18 @@ For information on using ESM3, see the [ESM3 README](https://github.com/Biohub/e
 
 ### Installation of ESM Package
 
-| Command | Contents |
-| :---- | :---- |
-| `pip install esm` | The models and the reference PyTorch path. No CUDA kernels. |
-| `pip install "esm[gpu]"` | Adds the xformers and cuequivariance kernels. Linux x86-64 only; a no-op elsewhere. |
-
-> **Changed in this release.** `pip install esm` no longer installs the
-> cuequivariance backend. Install `"esm[gpu]"` to keep it.
-
-Notes:
+```
+pip install esm
+```
 
-- Kernels are optional at runtime. The default install runs everywhere, it is only slower on a GPU.
-- Triton is not one of the extras. Torch requires it on Linux, so the `"fused"` backend already works in the default install.
+This installs the models, the reference PyTorch path, and Triton on Linux x86-64, so the ESMFold2 `"fused"` backend works without extra packages.
 
 #### Install size
 
-| Install | Size |
-| :---- | :---- |
-| `pip install esm` | about 5.6 GB |
-| `pip install "esm[gpu]"` | about 5.8 GB |
-
-The two are nearly the same size because torch on PyPI is always the CUDA
-build, about 4.6 GB of either total. The extra adds roughly 0.2 GB on top.
+`pip install esm` is about 5.6 GB, of which about 4.6 GB is the CUDA build of torch that PyPI always serves.
 
 Installing torch from its CPU-only index replaces that 4.6 GB CUDA build with
-an 0.8 GB CPU build, which is where the real saving is:
+an 0.8 GB CPU build:
 
 ```
 pip install esm --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
@@ -354,7 +341,7 @@ For high quality structure prediction, see the [maximum-accuracy folding SDK gui
 
 ## Batch API
 
-The Biohub Batch API processes larger or longer-running workloads in the background and offers higher limits. See `ForgeBatchClient` and the [Batch API guide](cookbook/snippets/batch.md) for examples of submission, progress checks, result retrieval, cache control, and cancellation. Queued tasks expire after 24 hours; results are retained for 31 days after submission.
+The Biohub Batch API processes larger or longer-running workloads in the background and offers higher limits, including folding inputs up to 4096 tokens rather than the Sync API's 768. See `ForgeBatchClient` and the [Batch API guide](cookbook/snippets/batch.md) for examples of submission, progress checks, result retrieval, cache control, and cancellation. Queued tasks expire after 24 hours; results are retained for 31 days after submission.
 
 ## Parallel Job Execution
 <a name="parallel-job-execution"></a>

cookbook/foldcp/README.md

@@ -24,7 +24,7 @@ Start from the NVIDIA PyTorch container, then install the dependencies:
 docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:26.03-py3 bash -l
 
 # ESM plus the Triton kernels
-pip install "esm[gpu] @ git+https://github.com/Biohub/esm.git@main"
+pip install "esm @ git+https://github.com/Biohub/esm.git@main"
 ```
 
 Context parallelism additionally needs `transformer-engine`, which is not

cookbook/foldcp/fold-cp.md

@@ -20,7 +20,7 @@ pins, so take the prebuilt wheel from the
 [EvolutionaryScale wheels release](https://github.com/evolutionaryscale/wheels/releases).
 
 ```bash
-pip install "esm[gpu] @ git+https://github.com/Biohub/esm.git@main"
+pip install "esm @ git+https://github.com/Biohub/esm.git@main"
 ```
 
 Use [esmfold2-cp.py](esmfold2-cp.py) as the CP version of the single-GPU

cookbook/foldcp/single-gpu.md

@@ -75,11 +75,7 @@ The "fused" backend uses Triton, a Python-based language for writing custom GPU
 
 ### Dependencies
 
-The "fused" backend needs Triton 3. It's GPU-only (Triton JIT-compiles to PTX) and **inference-only** (falls back to reference under autograd).
-
-```bash
-pip install "esm[gpu]"  # brings triton>=3,<4
-```
+The "fused" backend needs Triton 3, which `pip install esm` installs on Linux x86-64. It's GPU-only (Triton JIT-compiles to PTX) and **inference-only** (falls back to reference under autograd).
 
 > If Triton is not importable, `set_kernel_backend("fused")` silently runs the reference path. Check that Triton can be imported if `"fused"` is unexpectedly slow.
 

cookbook/snippets/batch.md

@@ -12,6 +12,8 @@ Use the Biohub Batch API to process larger or longer-running workloads in the ba
 | `fold_all_atom` | Molecular-complex prediction with ESMFold2 | [Input guide](../../README.md#running-esmfold2-through-the-biohub-platform) |
 | `fold_max_accuracy` | Maximum-accuracy molecular-complex prediction | [SDK guide](fold_max_accuracy.md) |
 
+Every fold service accepts inputs up to 4096 tokens, against 768 over the Sync API.
+
 We will be launching more models/services soon!
 
 ## Submit, wait, and retrieve results

cookbook/snippets/esmc.py

@@ -70,8 +70,8 @@ def raw_forward(model: EsmcForMaskedLM):
     tokenizer = EsmcTokenizer()
     encoded = tokenizer(sequences, return_tensors="pt", padding=True)
     output = model(
-        input_ids=encoded["input_ids"],
-        attention_mask=encoded["attention_mask"],
+        input_ids=encoded["input_ids"].to(model.device),
+        attention_mask=encoded["attention_mask"].to(model.device),
         output_hidden_states=True,
     )
     logits, embeddings, hiddens = (

esm/models/esmfold2/layers.py

@@ -726,8 +726,40 @@ def __init__(
             else _gated_residual_raw
         )
 
-    def forward(self, x: Tensor, c_l: Tensor, attention_params: tuple) -> Tensor:
-        mod = self.adaln_modulation(c_l)
+    def _adaln_modulation_cached(
+        self, c_l: Tensor, inference_cache: dict[str, Any] | None
+    ) -> Tensor:
+        """``adaln_modulation(c_l)``, computed once per fold instead of per step.
+
+        ``c_l`` is the diffusion conditioning, which ``inference_cache`` already
+        holds fixed for the whole fold, so every diffusion step recomputes the
+        identical SiLU + Linear. One slot per block instance, in the caller's
+        per-fold cache, so nothing survives ``fold()``.
+
+        Keyed on the *identity* of ``c_l``: a different tensor object recomputes.
+        That makes the cache correct even where the conditioning is not actually
+        step-invariant -- it only stops paying off. In-place mutation of a cached
+        ``c_l`` would go unnoticed; nothing in the diffusion path does that.
+        """
+        if inference_cache is None or torch.is_grad_enabled() or c_l.requires_grad:
+            return self.adaln_modulation(c_l)
+        slot: dict[str, Any] = inference_cache.setdefault("atom_adaln", {}).setdefault(
+            self, {}
+        )
+        if slot.get("c_l") is not c_l:
+            slot.clear()
+            slot["c_l"] = c_l
+            slot["mod"] = self.adaln_modulation(c_l)
+        return slot["mod"]
+
+    def forward(
+        self,
+        x: Tensor,
+        c_l: Tensor,
+        attention_params: tuple,
+        inference_cache: dict[str, Any] | None = None,
+    ) -> Tensor:
+        mod = self._adaln_modulation_cached(c_l, inference_cache)
         if mod.dim() == 2:
             mod = mod.unsqueeze(1)
         shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod.chunk(6, dim=-1)
@@ -796,10 +828,11 @@ def forward(
         c_l: Tensor,
         attention_params: tuple,
         return_intermediates: bool = False,
+        inference_cache: dict[str, Any] | None = None,
     ) -> Tensor | tuple[Tensor, list[Tensor]]:
         intermediates: list[Tensor] = []
         for block in self.blocks:
-            q_l = block(q_l, c_l, attention_params)
+            q_l = block(q_l, c_l, attention_params, inference_cache=inference_cache)
             if return_intermediates:
                 intermediates.append(q_l)
         if return_intermediates:
@@ -941,13 +974,24 @@ def forward(
             r_to_q = self.coords_linear(r_input)
             q = q + r_to_q
 
-        c = c.repeat_interleave(num_diffusion_samples, 0)
+        # repeat_interleave allocates a fresh tensor of identical values on every
+        # diffusion step. Cache the expanded conditioning next to c_base so the
+        # blocks downstream see one object per fold and can key caches on it.
+        if layer_cache is not None:
+            c_exp = layer_cache.get("c_expanded")
+            if c_exp is None:
+                c_exp = c.repeat_interleave(num_diffusion_samples, 0)
+                layer_cache["c_expanded"] = c_exp
+            c = c_exp
+        else:
+            c = c.repeat_interleave(num_diffusion_samples, 0)
 
         result = self.atom_transformer(
             q_l=q,
             c_l=c,
             attention_params=attention_params,
             return_intermediates=return_intermediates,
+            inference_cache=inference_cache,
         )
         if return_intermediates:
             q, intermediates = result
@@ -1018,6 +1062,7 @@ def forward(
         atom_attention_mask: Tensor,
         num_diffusion_samples: int = 1,
         return_intermediates: bool = False,
+        inference_cache: dict[str, Any] | None = None,
     ) -> tuple[Tensor, list[Tensor]]:
         """Returns (r_update, intermediates)."""
         atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0)
@@ -1030,6 +1075,7 @@ def forward(
             c_l=c_l,
             attention_params=p_lm,
             return_intermediates=return_intermediates,
+            inference_cache=inference_cache,
         )
         if return_intermediates:
             q_l, intermediates = result
@@ -1686,6 +1732,7 @@ def forward(
             atom_attention_mask=ref_mask,
             num_diffusion_samples=num_diffusion_samples,
             return_intermediates=return_atom_repr,
+            inference_cache=inference_cache,
         )
 
         # Step 8: compute denoised output

esm/models/esmfold2/output.py

@@ -26,9 +26,8 @@ def build_molecular_complex_from_features(
 ) -> MolecularComplex:
     """Construct a MolecularComplex from feature-dict tensors and chain metadata.
 
-    Non-polymer chains (ligands) collapse all per-atom tokens into a single
-    residue token whose pLDDT is the per-token average and whose hetero flag
-    is True.
+    Tokens are grouped into residues by ``residue_index`` (one per CCD component
+    for ligands), with per-residue pLDDT averaged over tokens.
     """
     mask_np = atom_mask.bool().cpu().numpy()
     coords_np = coords.float().cpu().numpy()
@@ -54,34 +53,7 @@ def build_molecular_complex_from_features(
         is_nonpolymer = ci.mol_type == MOL_TYPE_NONPOLYMER
         entity_info[ci.entity_id] = "non-polymer" if is_nonpolymer else "polymer"
 
-        if is_nonpolymer:
-            residue_name = ci.tokens[0].residue_name if ci.tokens else "LIG"
-            sequence_tokens.append(residue_name)
-            chain_ids_per_token.append(ci.asym_id)
-            avg_plddt = (
-                float(np.mean([plddt_np[ti.token_index] for ti in ci.tokens]))
-                if ci.tokens
-                else 0.0
-            )
-            confidence.append(avg_plddt)
-            token_atom_start = out_atom_cursor
-            for ti in ci.tokens:
-                for atom_idx in range(ti.atom_start, ti.atom_start + ti.atom_count):
-                    if not mask_np[atom_idx]:
-                        continue
-                    flat_positions.append(coords_np[atom_idx].tolist())
-                    flat_elements.append(get_element_symbol(int(elements_np[atom_idx])))
-                    chars = name_chars_np[atom_idx]
-                    name = "".join(
-                        chr(int(c) + 32) for c in chars if int(c) != 0
-                    ).strip()
-                    flat_names.append(name)
-                    flat_hetero.append(True)
-                    out_atom_cursor += 1
-            token_to_atoms.append([token_atom_start, out_atom_cursor])
-            continue
-
-        # Atom-tokenized modified residues (HYP, MSE, ...) span multiple
+        # Atom-tokenized residues (HYP, MSE, ligands, ...) span multiple
         # tokens per residue; collapse them back to one mmCIF residue.
         for _residue_index, ti_iter in groupby(
             ci.tokens, key=lambda t: t.residue_index
@@ -104,7 +76,7 @@ def build_molecular_complex_from_features(
                         chr(int(c) + 32) for c in chars if int(c) != 0
                     ).strip()
                     flat_names.append(name)
-                    flat_hetero.append(False)
+                    flat_hetero.append(is_nonpolymer)
                     out_atom_cursor += 1
             token_to_atoms.append([token_atom_start, out_atom_cursor])
 

esm/models/esmfold2/processor.py

@@ -183,6 +183,11 @@ def clean_esmfold2_input(input: StructurePredictionInput) -> StructurePrediction
     )
 
 
+def _on_cpu(t: torch.Tensor | None, i: int = 0) -> torch.Tensor | None:
+    """``t[i]`` detached on the host, or ``None`` when the output is absent."""
+    return t[i].detach().cpu() if t is not None else None
+
+
 class ESMFold2InputBuilder:
     def __init__(self, ccd_cache: Path | None = None):
         load_ccd(ccd_cache)
@@ -283,6 +288,15 @@ def decode(
         emb_seq_t = output.get("output_embedding_sequence")
         emb_pair_pooled_t = output.get("output_embedding_pair_pooled")
 
+        # Sample-invariant outputs: indexing [0] inside the loop re-ran the DtoH
+        # copy once per diffusion sample. The distogram is L*L*64 fp32, so at
+        # L=1000, Bm=5 that was 5 x 256 MB of pageable copies.
+        distogram_c = _on_cpu(distogram_t)
+        residue_index_c = _on_cpu(residue_index_t)
+        entity_id_c = _on_cpu(entity_id_t)
+        emb_seq_c = _on_cpu(emb_seq_t)
+        emb_pair_pooled_c = _on_cpu(emb_pair_pooled_t)
+
         results: list[MolecularComplexResult] = []
         for i in range(Bm):
             mc = build_molecular_complex_from_features(
@@ -300,35 +314,13 @@ def decode(
                     plddt=plddts[i].detach().cpu(),
                     ptm=float(ptm_t[i].item()) if ptm_t is not None else None,
                     iptm=float(iptm_t[i].item()) if iptm_t is not None else None,
-                    pae=pae_t[i].detach().cpu() if pae_t is not None else None,
-                    distogram=(
-                        distogram_t[0].detach().cpu()
-                        if distogram_t is not None
-                        else None
-                    ),
-                    pair_chains_iptm=(
-                        pair_chains_t[i].detach().cpu()
-                        if pair_chains_t is not None
-                        else None
-                    ),
-                    residue_index=(
-                        residue_index_t[0].detach().cpu()
-                        if residue_index_t is not None
-                        else None
-                    ),
-                    entity_id=(
-                        entity_id_t[0].detach().cpu()
-                        if entity_id_t is not None
-                        else None
-                    ),
-                    output_embedding_sequence=(
-                        emb_seq_t[0].detach().cpu() if emb_seq_t is not None else None
-                    ),
-                    output_embedding_pair_pooled=(
-                        emb_pair_pooled_t[0].detach().cpu()
-                        if emb_pair_pooled_t is not None
-                        else None
-                    ),
+                    pae=_on_cpu(pae_t, i),
+                    distogram=distogram_c,
+                    pair_chains_iptm=_on_cpu(pair_chains_t, i),
+                    residue_index=residue_index_c,
+                    entity_id=entity_id_c,
+                    output_embedding_sequence=emb_seq_c,
+                    output_embedding_pair_pooled=emb_pair_pooled_c,
                 )
             )
 

esm/sdk/__init__.py

@@ -17,15 +17,15 @@ def client(
     model="esm3-sm-open-v1",
     url="https://biohub.ai",
     token=os.environ.get("ESM_API_KEY", ""),
-    request_timeout=None,
+    request_timeout: int | None = 125,
 ) -> ESM3InferenceClient:
     """
     Args:
         model: Name of the esm3 model to use.
         url: URL of a Forge/Biohub Platform server.
         token: User's API token.
         request_timeout: Amount of time to wait for a request to finish.
-            Default is wait indefinitely.
+            None waits indefinitely, which can hang on a dropped connection.
     """
     if not model.startswith("esm3"):
         raise ValueError(f"Invalid model name: {model}")
@@ -36,15 +36,15 @@ def esmc_client(
     model="esmc-300m-2024-12",
     url="https://biohub.ai",
     token=os.environ.get("ESM_API_KEY", ""),
-    request_timeout=None,
+    request_timeout: int | None = 125,
 ) -> ESMCInferenceClient:
     """
     Args:
         model: Name of the esmc model to use.
         url: URL of a Forge/Biohub Platform server.
         token: User's API token.
         request_timeout: Amount of time to wait for a request to finish.
-            Default is wait indefinitely.
+            None waits indefinitely, which can hang on a dropped connection.
     """
     if not model.startswith("esmc"):
         raise ValueError(f"Invalid model name: {model}")
@@ -55,15 +55,15 @@ def esmfold2_client(
     model="esmfold2-fast-2026-05",
     url="https://biohub.ai",
     token=os.environ.get("ESM_API_KEY", ""),
-    request_timeout=None,
+    request_timeout: int | None = 125,
 ) -> SequenceStructureForgeInferenceClient:
     """
     Args:
         model: Name of the ESMFold2 model to use.
         url: URL of a Forge/Biohub Platform server.
         token: User's API token.
         request_timeout: Amount of time to wait for a request to finish.
-            Default is wait indefinitely.
+            None waits indefinitely, which can hang on a dropped connection.
     """
     if not model.startswith("esmfold2"):
         raise ValueError(f"Invalid model name: {model}")