NVIDIA / NVIDIA/TensorRT-Model-Connect

[Enterprise] Add end-to-end ESMFold2 support

Open
#1,113 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement Enterprise
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.compile baselines, 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.