NVIDIA / NVIDIA/TensorRT-Model-Connect
[Enterprise] Add end-to-end ESMFold2 support
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 254
- Forks
- 58
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 235
Description
Parent epic: #998
Summary
Add an end-to-end ESMFold2 path through TensorRT Model Connect for protein structure prediction.
The implementation must preserve ESMFold2-specific language-model conditioning and folding semantics rather than treating it as a variant of the Boltz-2 or OpenFold3 graph.
Tasks
- Pin the exact ESMFold2 source repository, source revision, checkpoint revision, configuration, tokenizer, and applicable licenses.
- Define supported amino-acid sequence, optional MSA, masking, feature, and batching inputs.
- Define the standard structure, confidence, and metadata output contracts.
- Implement the ESMC-conditioned model graph, folding trunk, diffusion path, and required model-owned transformations.
- Package tokenization, model configuration, inference components, and output metadata into a reproducible Model Connect bundle.
- Produce a standard structure artifact such as PDB or mmCIF through a native runtime example.
- Validate structures and confidence outputs against the pinned reference implementation using documented metrics and tolerances.
- Establish an initial BF16 qualification profile on at least one documented supported NVIDIA GPU.
- Benchmark representative sequence lengths for warmed steady-state latency, throughput where applicable, and peak memory against aligned eager PyTorch and
torch.compilebaselines, excluding engine-build time. - Characterize maximum supported sequence length, batching limits, memory use, valid-output rate, and known limitations.
- Add model-owned unit, integration, qualification, and end-to-end tests.
- Document clean-environment setup, bundle generation, native execution, structure output, and result interpretation.
Acceptance criteria
A user can provide a supported protein input to a pinned ESMFold2 configuration and receive a valid structure artifact through a native TensorRT Model Connect bundle within the agreed quality tolerances.
The qualified BF16 profile records the exact GPU and software stack and includes reproducible quality, validity, latency, and memory evidence plus documented sequence-length and configuration limits.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with parent epic #998 and establish the pinned ESMFold2 reference, revisions, configuration, tokenizer, and licenses listed in the tasks. Then trace the existing TensorRT Model Connect model path and compare it with the reference implementation. Done means a reproducible native bundle produces validated structure and confidence outputs, with BF16 quality, latency, memory, and sequence-limit evidence documented.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- documentation, machine-learning, performance, testing-qa
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 28/100