Kernels:

---

Rocketknight1

/

esmfold2-trimul-kernel

like 0

# esmfold2-trimul-kernel

Fused inference Triton kernel for **ESMFold2's triangle multiplication** — the
O(N³) hotspot of the folding trunk (76 sites in the full model). Packaged as a
[`kernels`](https://github.com/huggingface/kernels)-library Hub kernel so
`transformers` can load it on demand; the pure-PyTorch block in transformers is the
fallback.

## What it fuses

One Triton kernel for the whole `EsmFold2TriangleMultiplicativeUpdate`:
`norm_start → gated dual-GEMM (sigmoid(x@Wg)·(x@Wp)) → triangular einsum
(bikd,bjkd→bijd) → norm_mix → proj_emit → output gate`, with the `delta` intermediate
never written to HBM. bf16 in/out, fp32 accumulation. Forward + backward present;
registered inference-only in transformers.

## How transformers uses it

`EsmFold2TriangleMultiplicativeUpdate` is decorated `@use_kernel_forward_from_hub(
"ESMFold2TriangleMultiplication")` and mapped to this repo in
`integrations/hub_kernels.py` (cuda, `Mode.INFERENCE`). The layer here
(`ESMFold2TriangleMultiplication`) reimplements that module's `forward(pair_grid,
visibility)`, reading its parameters (`norm_start`/`norm_mix`/`proj_bundle`/
`proj_emit`/`proj_gate`, plus `dim`/`flow`). **Keep this layer in sync with the in-tree module's
attribute names and forward signature** — that's the contract.

```python
import torch
from transformers import EsmFold2Model

# use_kernels=True swaps in this kernel for the 76 trimul sites (CUDA + inference).
model = EsmFold2Model.from_pretrained(
    "biohub/ESMFold2", dtype=torch.bfloat16, device_map="cuda", use_kernels=True
).eval()
out = model.infer_protein(seq)
```

## Layout

The package name must match `kernels`' repo-derived name
(`repo_id.split("/")[-1].replace("-", "_")`), i.e. **`esmfold2_trimul_kernel`** for the
repo `…/esmfold2-trimul-kernel`, and `build.toml`'s `[general] name` must match it too.
`kernels.get_kernel` loads from `build/torch-universal/` (not `torch-ext/`).

```
build.toml                                       # kernel-builder config (universal/Triton)
flake.nix                                        # kernel-builder entry (verify vs current version)
torch-ext/esmfold2_trimul_kernel/                # source (read by kernel-builder)
  __init__.py                                    # exports the layer + the functional entry
  layers.py                                      # ESMFold2TriangleMultiplication (the Hub layer)
  trimul_with_residual.py                        # kernel entrypoint
  fused_dual_gemm.py                             # helper: gated dual GEMM
  fused_ln_residual.py                           # helper: LN + transpose / residual-link epilogues
  trimul_einsum_triton.py                        # helper: batched triangular einsum
build/torch-universal/esmfold2_trimul_kernel/    # loaded by kernels.get_kernel (same files)
```

## Build & publish

The `build/torch-universal/` dir checked in here is a hand-built universal layout
(the Triton package copied in — no compile step), which is sufficient for
`kernels.get_kernel`. To regenerate it properly with kernel-builder:

```bash
nix build .#bundle      # or: kernel-builder build  (see kernel-builder docs)
```

Mapped from transformers in `integrations/hub_kernels.py` under the layer name
`ESMFold2TriangleMultiplication`. Currently `repo_id = Rocketknight1/esmfold2-trimul-kernel`
(testing); move it to a `kernels-community` org and update `repo_id` before merging.

Two Hub-side requirements, both enforced by `kernels` at load time:

1. **The repo must be of type `kernel`**, not `model`. `kernels` resolves revisions via
   `list_repo_refs(repo_type="kernel")`, and a plain model repo 404s there. Check with
   `HfApi().list_repo_refs(repo_id, repo_type="kernel")`.
2. **A version tag must exist** when the `LayerRepository` pins `version=N` (it currently
   pins `version=1`). Version resolution reads `vX.Y.Z` tags — compare
   `kernels-community/deformable-detr`, which carries `v0.0.1`/`v0.0.2`/`v0.0.3`. With no
   tags, resolution fails regardless of the code. Either push a `v1.0.0` tag, or switch
   `hub_kernels.py` to a pinned `revision=` and drop `version=`.

A failure in either surfaces as a `RepositoryNotFoundError` raised from
`LayerRepository.__str__` while formatting an unrelated message, which hides the real cause.

## Validation

Swapped into all 76 `EsmFold2TriangleMultiplicativeUpdate` instances of the real model
(`biohub/ESMFold2`, bf16, GPU), folds match the pure-PyTorch fallback within the
model's own non-determinism: ubiquitin 0.801 vs 0.799 pLDDT (Δ +0.002), GB1 0.849 vs
0.849, pTM identical. Standalone microbench (dim=128, B=1): 5–37× over the chunked
fp32 fallback, gap growing with N (`torch.compile` of the fallback only reaches
~1–7×).

## Follow-ups

- **Residual-optional entry.** The in-tree boundary is delta-only, so the layer
  passes `residual=zeros_like(pair)`, costing one `[B,N,N,C]` alloc+read per call.
  Adding a `residual=None` fast path (skip the in-kernel residual add) recovers that.
- **cuequivariance** provides the same op (`triangle_multiplicative_update`) as an
  alternative backend if a vendored-Triton kernel is undesirable.