Skip to content

[Enterprise] Add end-to-end ESMFold2 support #1113

Description

@yifeif-nv

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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    EnterpriseEnterprise related issuesenhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions