[OpenFold] Custom parts / AlphaFold

[blefaudeux]

🚀 Feature

Checkout the discussion here, OpenFold is missing a couple of custom/optimized parts in PyTorch and maybe that xformers could host some of these

Motivation

• Transformer related
• Custom parts, not handled out of the box by the PyTorch primitives in a fast enough fashion
• Possibly good candidates for Triton, Functorch or Sputnik, but would need more reading..

Pitch

Implement some of these primitives in xformers

Alternatives

Well, do nothing

Additional context

AlphaFold2 supplementary data: https://static-content.springer.com/esm/art%3A10.1038%2Fs41586-021-03819-2/MediaObjects/41586_2021_3819_MOESM1_ESM.pdf

[blefaudeux]

blocked by current triton (consecutive matmuls in a given kernel are not well supported), sorry about that @gahdritz, I'll revisit this over time