release batch API, max accuracy folding

README.md

@@ -42,6 +42,7 @@ For information on using ESM3, see the [ESM3 README](https://github.com/Biohub/e
 - [ESMC](#esmc)
 - [ESMC Sparse Autoencoders](#esmc-sparse-autoencoders)
 - [ESMFold2](#esmfold2)
+- [Batch API](#batch-api)
 - [Parallel Job Execution](#parallel-job-execution)
 - [Frontier-Safety](#frontier-safety)
 - [Licenses](#licenses)
@@ -50,7 +51,44 @@ For information on using ESM3, see the [ESM3 README](https://github.com/Biohub/e
 ## ESM on PyPI
 <a name="esm-on-pypi"></a>
 
-[`esm v3.4.post1`](https://github.com/Biohub/esm/releases/tag/v3.4.1.post1) is now available on PyPI with full support for ESMC and ESMFold2.
+[`esm v3.4.1.post1`](https://github.com/Biohub/esm/releases/tag/v3.4.1.post1) is now available on PyPI with full support for ESMC and ESMFold2.
+
+### 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:
+
+- 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.
+
+#### 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.
+
+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:
+
+```
+pip install esm --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
+```
+
+Worth doing if you have no NVIDIA GPU, are on Apple Silicon, or want smaller CI
+images and containers.
+
+`flash-attn` and `transformer-engine` are not installable from PyPI against our torch pin. Prebuilt wheels for both are published in the [EvolutionaryScale wheels release](https://github.com/evolutionaryscale/wheels/releases).
 
 ## ESM Package vs. Hugging Face Transformers Usage
 <a name="esm-package-vs-hugging-face-transformers-usage"></a>
@@ -150,7 +188,7 @@ For tutorials on how to use ESMC, see our [tutorials](https://github.com/Biohub/
 
 Sparse autoencoders (SAE) are an unsupervised method for decomposing representations of large transformer language models into interpretable units. We released SAEs trained on ESMC to reveal the interpretable units of functional organization that ESMC's world model has learned.
 
-The sparse autoencoder used in the Atlas and analyzed in the paper, `ESMC-6B-sae-layer60-k64-codebook16384`, is built on the ESMC 6B model. We also provide human-interpretable, agent-generated feature descriptions for this SAE's codebook.
+The sparse autoencoder used in the example below, `ESMC-6B-sae-k64-codebook16384`, is built on the ESMC 6B model and covers all layers. The SAE used in the Atlas and analyzed in the paper is `ESMC-6B-sae-layer60-k64-codebook16384`; the human-interpretable, agent-generated feature descriptions apply to its codebook.
 
 Codebase, model weights, and model variants for ESMC SAEs are available through [Hugging Face](https://huggingface.co/collections/biohub/esmc-saes-for-hidden-states-all-layers).
 
@@ -266,7 +304,7 @@ with open("1mht_pred.cif", "w") as f:
     f.write(result.complex.to_mmcif())
 ```
 
-> **AMD ROCm users:** use ROCm 6.4 with PyTorch 2.9 or newer.
+> **AMD ROCm users:** use ROCm 6.4 or newer. `esm` requires PyTorch 2.11.
 
 ### Running ESMFold2 Through the Biohub Platform
 
@@ -312,9 +350,15 @@ with open("result.cif", "w") as f:
 
 For tutorials on how to use ESMFold2, see our [tutorials](https://github.com/Biohub/esm/tree/main/cookbook/tutorials).
 
+For high quality structure prediction, see the [maximum-accuracy folding SDK guide](cookbook/snippets/fold_max_accuracy.md).
+
+## 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.
+
 ## Parallel Job Execution
 <a name="parallel-job-execution"></a>
-For jobs that require processing multiple inputs, the Parallel Executor sends parallel requests, executing them concurrently and efficiently while respecting rate limits and adapting to request latency. The following example provides an example of using the parallel_executor context manager when embedding sequences.
+Use `parallel_executor` to send multiple Sync API requests concurrently while respecting rate limits and adapting to request latency. For queued background work, use `ForgeBatchClient`. The following example uses `parallel_executor` to embed sequences through the Sync API.
 
 ```py
 from esm.sdk.forge import ESMCForgeInferenceClient
@@ -334,7 +378,7 @@ sequences = ["A", "AA", "AAA"]
 client = ESMCForgeInferenceClient(model="esmc-6b-2024-12", url="https://biohub.ai", token="<your API token>")
 
 # Usage Example:
-# To execute a batch job, wrap your function inside the batch executor context manager.
+# To process multiple inputs through the Sync API, wrap your function in parallel_executor.
 # Syntax:
 # with parallel_executor() as executor:
 # outputs = executor.execute_batch(user_func=<your_function>, **kwargs)

cookbook/foldcp/README.md

@@ -11,10 +11,9 @@ Tested on 1x and 4x H100 80GB GPUs.
 | Script | Input | GPUs | Result |
 | --- | --- | --- | --- |
 | `esmfold2-none.py` | 7ysz (1 protein chain ×2) | 1 | Reference single-GPU path |
-| `esmfold2-cueq.py` | 7ysz (1 protein chain ×2) | 1 | cuEquivariance backend |
 | `esmfold2-fused.py` | 7ysz (1 protein chain ×2) | 1 | Triton fused backend |
 | `esmfold2-cp.py` | 7ysz (1 protein chain ×2) | 4 | Same fold API with CP setup |
-| `cuequivariance_cp.py` | 5xgo (1 protein chain ×12, CL ligand) | 4 | Larger input via CP |
+| `large_input_cp.py` | 5xgo (1 protein chain ×12, CL ligand) | 4 | Larger input via CP |
 | `will_fail.py` | 5xgo (same as CP) | 1 | **OOMs** — too large for 1xH100 |
 
 ## Environment setup
@@ -24,27 +23,23 @@ Start from the NVIDIA PyTorch container, then install the dependencies:
 ```bash
 docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:26.03-py3 bash -l
 
-# Main ESM package
-pip install "git+https://github.com/Biohub/esm.git@main"
-# fused dependency
-pip install "esm[fused] @ git+https://github.com/Biohub/esm.git@main"
-# cuEquivariance dependency (cueq12 also exists)
-pip install "esm[cueq13] @ git+https://github.com/Biohub/esm.git@main"
-# fold-cp dependency for larger inputs
-pip install "esm[fold-cp] @ git+https://github.com/Biohub/esm.git@main"
+# ESM plus the Triton kernels
+pip install "esm[gpu] @ git+https://github.com/Biohub/esm.git@main"
 ```
 
+Context parallelism additionally needs `transformer-engine`, which is not
+installable from PyPI against our torch pin. Take the prebuilt wheel from the
+[EvolutionaryScale wheels release](https://github.com/evolutionaryscale/wheels/releases).
+
 ## Single-GPU accelerated backends
 
-[single-gpu.md](single-gpu.md) compares the three single-GPU backend choices:
-`None` for the pure PyTorch reference path, `"cuequivariance"` for NVIDIA
-cuEquivariance triangle-multiplication kernels, and `"fused"` for Triton kernels
-that fuse several ESMFold2 hot-path operations. These backends keep the same model
+[single-gpu.md](single-gpu.md) compares the two single-GPU backend choices:
+`None` for the pure PyTorch reference path and `"fused"` for Triton kernels
+that fuse several ESMFold2 hot-path operations. Both backends keep the same model
 weights and outputs while trading dependencies, warm-up behavior, and speed.
 
 ```bash
 python esmfold2-none.py
-python esmfold2-cueq.py
 python esmfold2-fused.py
 ```
 
@@ -62,5 +57,5 @@ torchrun --nproc-per-node=4 esmfold2-cp.py
 
 # Larger 5xgo example.
 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
-  torchrun --nproc-per-node=4 cuequivariance_cp.py
+  torchrun --nproc-per-node=4 large_input_cp.py
 ```

cookbook/foldcp/fold-cp.md

@@ -14,10 +14,13 @@ proteins than fit on a single GPU. Wrapping results in the same object with a fe
 
 ## Quickstart
 
-Install the fold-cp dependency
+Install ESM, then `transformer-engine`. PyPI only ships it as an sdist that needs
+a CUDA toolchain, and its core wants a cuBLAS symbol newer than the one torch
+pins, so take the prebuilt wheel from the
+[EvolutionaryScale wheels release](https://github.com/evolutionaryscale/wheels/releases).
 
 ```bash
-pip install "esm[fold-cp] @ git+https://github.com/Biohub/esm.git@main"
+pip install "esm[gpu] @ git+https://github.com/Biohub/esm.git@main"
 ```
 
 Use [esmfold2-cp.py](esmfold2-cp.py) as the CP version of the single-GPU
@@ -43,7 +46,7 @@ unchanged.
 
 - Python environment:
   - This ESM package
-  - transformer-engine 2 (installed by `esm[fold-cp]` dependency)
+  - transformer-engine 2, from the EvolutionaryScale wheels release
 - Use a perfect-square number of GPUs: 1, 4, 9, 16, ...
 - Keep `num_diffusion_samples=1` in `ESMFold2InputBuilder().fold(...)`. The CP
   diffusion path currently expects one diffusion sample.
@@ -111,7 +114,7 @@ Two differences are expected, both small:
 - **Rounding** differs at about 1e-3, because numbers are summed across GPUs in a
   different order.
 
-Single-GPU kernel backends such as `"fused"` and `"cuequivariance"` are ignored by the CP-wrapped stages; those stages use the distributed CP implementations instead.
+The single-GPU `"fused"` kernel backend is ignored by the CP-wrapped stages; those stages use the distributed CP implementations instead.
 
 ## Larger example
 
@@ -125,7 +128,7 @@ Context parallelism sharding will split the memory and allow this protein to be
 Launch with torchrun on a perfect-square number of GPUs:
 
 ```bash
-PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True torchrun --nproc-per-node=4 cuequivariance_cp.py
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True torchrun --nproc-per-node=4 large_input_cp.py
 ```
 
 ## How context parallelism works

cookbook/foldcp/cuequivariance_cp.py → cookbook/foldcp/large_input_cp.py

@@ -30,7 +30,7 @@
     ]
 )
 
-# torchrun entrypoint:  torchrun --nproc_per_node=4 cuequivariance_cp.py
+# torchrun entrypoint:  torchrun --nproc_per_node=4 large_input_cp.py
 # torchrun sets LOCAL_RANK / WORLD_SIZE / RANK / MASTER_ADDR / MASTER_PORT.
 # world_size must be a perfect square (the CP grid is n×n).
 local_rank = int(os.environ["LOCAL_RANK"])
@@ -65,7 +65,7 @@
 # Fused kernels accelerate the modules *outside* the CP region (notably the
 # diffusion sampler / structure head). The CP-wrapped trunk and MSA encoder
 # no-op this call, so it never touches the distributed ring math.
-model.set_kernel_backend("cuequivariance")
+model.set_kernel_backend("fused")
 
 start = time()
 with torch.inference_mode():