[vshah1016]

Summary

This PR makes ESM3.logits() usable inside a CUDA Graph. The current implementation creates its default confidence scalar through a CPU-backed operation; the resulting host-to-device copy fails during capture. Creating that scalar directly on the target device removes the capture failure.

The PR also converts only the outputs requested by LogitsConfig to FP32, adds regression tests, and provides a benchmark for repeated, fixed-shape inference. The cookbook includes an optional strict SwiGLU compilation recipe and records the alternatives we tested.

On one B200, the original 76-residue ubiquitin experiment measured 32.81 ms eager versus 6.269 ms with fresh input copy and graph replay, about 5.23× faster. A separate experiment reduced graph latency by a further 3.69–5.61% across three workloads by compiling only SwiGLU. Applications must explicitly capture and replay a graph to obtain these gains; ordinary calls remain eager.

1. Changes in this PR

Area	Previous behavior	Change and resulting behavior
Confidence scalar in logits()	torch.tensor(1.0, device=input.device) introduces a CPU-to-GPU transfer that fails during capture.	Use torch.ones((), device=input.device). The scalar retains its value, shape, and default floating dtype while being created on the target device.
Output conversion	Every output is converted to FP32 before unrequested tracks are discarded.	Convert each output only when requested. Requested outputs retain their device and FP32 dtype; unrequested outputs remain None.
Benchmark and cookbook	No accompanying recipe for this capture path.	Add changing-input eager/graph comparisons, output checks, raw timing samples, and optional compiler/backend experiments.
Setup instructions	The contribution guide names Python 3.10, a nonexistent requirements file, and renamed notebooks.	Use Python >=3.12, correct the paths, explain notebook token prompts, and provide an isolated performance environment.

The output-selection change still computes every model head. Both the model's forward hooks and output-head hooks continue to receive the complete raw outputs. The library's default precision, attention backend, and compilation settings are unchanged.

Files to review

Files	Purpose
esm/models/esm3.py	The scalar construction and conditional output conversion.
esm3_cuda_graph_test.py, esm3_logits_selection_test.py	Capture, changed-input replay, scalar semantics, output selection, dtypes, and hooks.
esm3_cuda_graph.py	Public-API benchmark and optional experiments.
Cookbook guide, recorded results	Usage, measurement method, results, and limitations.
CONTRIBUTIONS.md, README.md	Setup corrections and a link to the recipe.

2. Correctness and API behavior

The two regression modules contain 24 cases: 20 CPU cases and four CUDA cases. They cover:

• The confidence scalar under FP32 and FP64 default dtypes.
• No requested outputs, all outputs, and each individual output, using FP32 and BF16 raw results.
• The existing output-head and model forward hooks, including a hook that replaces raw outputs.
• CUDA Graph replay with changed sequences and, where present, changed coordinates, using FP32 and BF16 model weights. These CUDA cases compare replay with eager execution using torch.testing.assert_close.

Full-model benchmark checks provide separate evidence. For the accepted graph and strict-SwiGLU configurations, every requested tensor passed torch.equal against the appropriate eager reference across four changing inputs. The SwiGLU candidates also matched the original, uncompiled BF16 model. These are exact returned-value comparisons on the recorded fixtures; they do not establish equivalence for every possible input or biological task.

Applications own the graph's static input and output buffers. They must warm the model before capture and copy new inputs into the captured buffers. Returned graph outputs reuse storage, so an application must consume or clone them before the next replay. A different shape, input-track set, output configuration, or model configuration requires another warmed capture.

3. Performance results

A. Enabling CUDA Graph execution

The original capture experiment used one B200, batch size 1, and 76-residue ubiquitin encoded to 78 tokens. All six output heads and embeddings were requested.

Execution path	Median latency
Eager with fresh pinned-host input, control before / after	32.811 / 32.821 ms
Graph replay with fixed input, control before / after	6.2618 / 6.2620 ms
Graph replay with fresh device-input copy	6.2709 ms
Graph replay with fresh pinned-host input copy	6.2695 ms

The fresh-host-input path is 5.23× faster than the mean of the two eager control medians. Its increment over fixed-input replay is 7.56 microseconds for the 624-byte sequence input, approximately 0.12% of the complete graph call. That increment is a difference of measured medians, not a standalone transfer-bandwidth measurement.

The unmodified capture failure and patched timing/correctness record are retained. This comparison measures the execution path enabled by the patch; it does not claim that replacing the scalar alone accelerates eager inference by 5.23×.

B. Optional SwiGLU compilation inside the graph

Graph replay reduces CPU launch overhead, but the captured kernels still execute individually. A node-level Nsight Systems trace showed many small elementwise, conversion, and copy kernels. We therefore tested compiling only each block's SwiGLU activation, using numerics="strict" and disabling Inductor's own CUDA Graph handling.

The fresh-host rerun measured the following, with a graph control before and after each candidate:

Workload	Graph control before / after	Graph with compiled SwiGLU	Lower median latency
Batch 1, 76 residues	6.226 / 6.227 ms	5.997 ms	3.69%
Batch 8, 76 residues	9.203 / 9.204 ms	8.687 ms	5.61%
Batch 1, 512 synthetic residues	11.073 / 11.072 ms	10.619 ms	4.09%

All seven returned tensors passed exact comparisons across four input variants for each workload. A separate profile recorded 2,386 → 2,338 kernels per replay: 48 SiLU/multiply pairs became 48 fused kernels. Normalization, attention, and GEMM operations were left unchanged. Fresh-run samples · Profiler record

The original series measured reductions of 3.64%, 5.67%, and 4.11%. The fresh run is a replication by the same team on another host. It is separate from the capture experiment above; their timings are not combined into a new end-to-end speedup.

The first candidate call took 1.96–4.26 seconds in the fresh run, including compilation-related startup. This cost must be amortized over repeated calls. The original batch-1 series also contains a 14.482 ms outlier, retained in the data. These results support a median-latency improvement; they do not establish improved tail latency.

C. Selective output conversion

Sequence-only and embeddings-only requests each remove six FP32 conversion copies. Full-model comparisons passed exact output checks at batch 1 and 8 for ubiquitin and batch 1 for the 512-residue synthetic fixture.

The measured graph-latency difference was only about 0–0.2%, and whole-call peak allocation did not decrease in the batch-1 ubiquitin measurement. This change removes unnecessary work, but the measurements do not support a meaningful latency or peak-memory improvement. Selection measurements

4. Measurement method and reproduction

Recorded environment

Component	Performance experiments
Hardware	One NVIDIA B200 per comparison, on AWS p6-b200.48xlarge hosts
Driver	595.112.07
PyTorch / CUDA runtime	2.15.0+cu132 / 13.2
Triton / CPU threads	3.9.0 / OMP_NUM_THREADS=4
Model	ESM3 Open, checkpoint 47f0545b2b6daf26a93439a3cd610f4f7f3d5478
Profiling	Nsight Systems 2026.1.1 with node-level graph tracing; a separate Nsight Compute sample of one GEMM

The primary input is ubiquitin with four sequence variants: unmasked and three mask patterns. Batch rows cycle through the variants. Longer fixtures repeat ubiquitin before masking; they are synthetic shape tests.

The capture and reported SwiGLU experiments each use 20 warmups and 100 timed calls per path, with GPU synchronization after each call. Controls bracket the candidates. The tables report medians; reductions use the mean of the before/after control medians as the reference.

For the headline fresh-input and SwiGLU comparisons, timing includes the pinned-host input copy, model execution or graph replay, and synchronization. Section 3A labels the additional fixed-input and device-copy controls separately. Timing excludes tokenization, host allocation/pinning, capture, compilation, and copying outputs back to the CPU. Profiling is separate from the reported uninstrumented latency measurements. Preliminary attention/quantization runs whose later controls were disturbed by profiling are retained in the archive but excluded from the reported comparisons.

Consecutive calls measure variation within a run. They are not independent experimental trials, and no independent-run confidence interval was estimated. The measurements establish behavior on the tested B200 fixtures; biological quality, B300 performance, and a comparison against Hopper remain untested.

Running the public recipe

Start from this PR's checkout with Python >=3.12 and a working CUDA environment. The cookbook setup installs the measured performance stack and explains how to pin the checkpoint. Keep this environment separate from the repository's locked test environment.

After that setup, run from the checkout root:

export CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4
PYTHONPATH=. python cookbook/local/esm3_cuda_graph.py --output graph.json
PYTHONPATH=. python cookbook/local/esm3_cuda_graph.py --compile-swiglu --output swiglu-b1.json
PYTHONPATH=. python cookbook/local/esm3_cuda_graph.py --compile-swiglu --batch-size 8 --output swiglu-b8.json
PYTHONPATH=. python cookbook/local/esm3_cuda_graph.py --compile-swiglu --length 512 --output swiglu-l512.json

The JSON records source/checkpoint identity, output comparisons, setup costs, and individual timing samples. For the default and strict-SwiGLU paths, failing exact output checks produces a nonzero exit. Approximate experimental configurations have separately reported tolerance checks.

The complete reproduction guide covers the pinned GPU-tested revision, locked tests, performance environment, and FA2 source build. Its portable command sequence was assembled from executed scripts and has not itself been rerun end to end on another GPU host.

5. Validation status

CPU CI and formatting/typing were rerun at the current head, 34bbb9b. The final GPU validation and fresh-host benchmark rerun used d285dde; the later commit changes documentation only. Model code, regression tests, benchmark code/results, and dependency files are byte-identical between those two revisions. Full revision identifiers are in the appendix.

Check	Environment / revision	Result	Evidence
CPU test matrix	Locked Ubuntu, Python 3.12 / 3.13 / 3.14; 34bbb9b	592 passed, 3 skipped, 7 xfailed per version	Actions run
Formatting and typing	Full native checks, including whole-project ty; 34bbb9b	Passed	Job log
ESM3 regression modules	B200 host; Torch 2.11/CUDA 13.0 and Torch 2.15/CUDA 13.2; d285dde	24 passed in each environment, including four CUDA cases	2.11 JUnit, 2.15 JUnit
Broader GPU test selection	Locked Torch 2.11 with an FA2 SM100 overlay; d285dde	59 passed, zero failures/errors/skips; 352 tests deselected by marker selection	JUnit, log
Public benchmark	Torch 2.15, B200; d285dde	Default, sequence-only, embeddings-only, and three SwiGLU shapes passed exact checks	Six records
Required local contribution scripts	Unmodified checkout, local checkpoint mode; d285dde	snippets/esm3.py, snippets/esmc.py, and local/raw_forwards.py passed on the B200 host	Run summary
Docker/Forge integration	Fork workflow at 34bbb9b; image installs the published package	Four failures: missing ESM_API_KEY, before inference	Actions run
Four tutorial notebooks	Contribution checklist	Not run	Contribution guide
Upstream CI	34bbb9b	Awaiting maintainer approval (action_required); no jobs executed	Upstream run

These suites have different scopes and overlap; their counts are not added into a single total. The fork workflow remains red because of the Forge failures. Passing the local checks does not complete the hosted integration or notebook checks.

Why nine GPU tests initially failed, and how the environment was repaired

The initial GPU selection reported 40 passed and 9 failed. The failures were ESMC/ESMFold2 paths entering FlashAttention2 with no kernel image: the repository's dev wheel targets SM80/90, while B200 is SM100.

We built official FA2 2.8.3.post2 source for SM100 against locked Torch 2.11/CUDA 13.0 in a separate overlay. Native cubins and a forward/backward probe were verified before rerunning the suite. The final 59-case selection retained every original 49-case identity, including the nine failures, and added ten cached-checkpoint cases. No ESM source or test assertions changed during this repair.

The repair changed both the FA2 version and binary architecture; it does not isolate architecture flags alone. The ESM3 performance measurements used a separate Torch 2.15 environment without FA2 installed, so their speedups cannot be attributed to this rebuild.

Original failures · Build command · Native architectures · Probe · Reproduction instructions

6. Other experiments and related work

The benchmark exposes the following alternatives for investigation. Their measurements explain why they were not adopted as default settings.

Candidate	Observed result	Decision
Broader transformer or FFN compilation	Faster, but failed comparison with the original BF16 model at rtol=atol=1e-2. Whole-forward fullgraph=True also encountered the tokenizer's Rust lookup.	Retain as an experiment; not an equivalent replacement for the original execution path.
Forced cuDNN attention	No gain: the default dispatcher already selected cuDNN on this stack.	Keep the existing dispatcher.
FlashAttention4	1.4–4.4% slower across the tested shapes, with changed outputs.	Retain diagnostic support and negative results.
MXFP8 FFNs	Uncompiled execution was 65.7–81.8% slower. Matched compiled comparisons were slower or effectively unchanged, with output drift.	Retain as an opt-in experiment; no default quantization change.

Detailed comparisons and raw records include the original-reference checks. Agreement between a candidate's graph and its own eager execution alone does not establish agreement with the original model.

Earlier batching and fused-AdamW measurements used existing APIs and PyTorch settings. They are documented in the historical experiments, separately from this PR's implementation and inference-latency claims.

Two related changes remain outside this branch: #417 addresses raw BF16 forward(), and #430 addresses padding when sequence_id is absent. These measurements include neither change. Equal-length batching here does not validate mixed-length padding; integrating #430 would also need a capture check. Public logits() already establishes CUDA BF16 autocast, while direct forward() calls with the open BF16 weights still require the appropriate context.

7. Review request

Please review the scalar and output-selection changes, their regression coverage, and whether the optional compilation recipe fits the cookbook. The upstream workflow also needs maintainer approval to run. Hosted integration and notebook validation remain the outstanding checks listed above.

Evidence appendix: revisions, raw data, and verification

Record	Revision
Upstream base	fddac936fc921445da2d099a3e4f7ae928717292
Final GPU-tested implementation	d285ddeead83631bbe0bc6aa7c879eb5a8946ccb
Current PR head, including setup-documentation corrections	34bbb9be4e5e4d2dd6104392abf247abef2f1a48
Separate, immutable experiment archive	6b90888707f77dba4815e73ad11db508fd70cfa0

The archive records the November 12–13, 2026 experiments and the validation state at publication. Its historical references to draft status and stale contribution instructions precede the current documentation commit; the current status is recorded above.

• Complete experiment report: issues, attempted solutions, results, and limitations.
• Reproduction guide: pinned environments and validation commands.
• Evidence index: the 176 archived artifacts.
• Manifest: original and published SHA-256 hashes and sanitization notes.
• Offline verifier: checks integrity, timing arithmetic, and test outcomes without a GPU or network.

To check the archive, run python3 verify.py from its experiments/esm3-blackwell directory. The archive is separate from the implementation branch; historical probes and logs are evidence, not additional library changes.

Contributed by @CumulusLabsIO (Cumulus Labs).