esmfold2-trimul: support torch.compile, upload to biohub

esmfold2-trimul/README.md

@@ -34,8 +34,14 @@ registers it inference-only.
 ## 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`).
+`@use_kernel_forward_from_hub("EsmFold2TriangleMultiplication")` and mapped to this
+repo's `ESMFold2TriangleMultiplication` layer in `integrations/hub_kernels.py` (cuda,
+inference).
+
+The layer sets `can_torch_compile = True`: the inference path traces with
+`torch.compile(fullgraph=True)` (the Triton kernels are captured directly, no graph
+break), and the sequence length stays symbolic under dynamic shapes, so new lengths do
+not recompile.
 
 ```python
 import torch
@@ -69,9 +75,8 @@ output past their ends; the LayerNorm kernels additionally index channels with
 raise `ValueError` instead of silently returning nondeterministic garbage — see
 "Follow-ups" for the real fix.
 
-Sequence length `L` is unconstrained (the row dim is masked everywhere), except that
-`B == L == 1` makes `M == 1`, which Triton specializes to a constexpr and the LayerNorm
-kernel's `M.to(tl.int64)` then rejects.
+Sequence length `L` is unconstrained (the row dim is masked everywhere), including
+`B == L == 1`.
 
 ## Validation
 
@@ -86,7 +91,8 @@ growing with N (`torch.compile` of the fallback only reaches ~1–7×).
 `tests/` checks the fused path against an fp32 PyTorch reference across both flow
 directions, supported channel counts, mask/dropout-mask combinations and a range of
 sequence lengths, and asserts repeated calls are bit-identical (drift would indicate an
-out-of-bounds read).
+out-of-bounds read). It also checks that `torch.compile(fullgraph=True)` matches eager
+and that varying the sequence length under `dynamic=True` never recompiles.
 
 ## Follow-ups
 

esmfold2-trimul/build.toml

@@ -6,7 +6,7 @@ license = "Apache-2.0"
 backends = ["cuda"]
 
 [general.hub]
-repo-id = "kernels-community/esmfold2-trimul"
+repo-id = "biohub/esmfold2-trimul"
 
 [torch-noarch]
 

esmfold2-trimul/tests/test_esmfold2_trimul.py

@@ -2,7 +2,9 @@
 import torch
 import torch.nn.functional as F
 
-import esmfold2_trimul
+import kernels
+
+esmfold2_trimul = kernels.get_kernel("biohub/esmfold2-trimul", version=1)
 
 
 # Channel counts the fused GEMMs support: powers of two >= 64. ESMFold2 uses 128.
@@ -118,12 +120,12 @@ def test_masks(use_mask, use_drop_mask):
 
 @pytest.mark.kernels_ci
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
-@pytest.mark.parametrize("L", [2, 7, 33, 64])
+@pytest.mark.parametrize("L", [1, 2, 7, 33, 64])
 def test_sequence_lengths(L):
     """L is unconstrained: the row dim is masked in every kernel.
 
-    B == L == 1 is excluded: it makes M == 1, which Triton specializes to a
-    constexpr and the LayerNorm kernel's ``M.to(tl.int64)`` then rejects.
+    B == L == 1 makes M == 1, which Triton specializes to a constexpr; the
+    LayerNorm kernels must still accept it.
     """
     _assert_matches_reference(1, L, 64, "outgoing", use_mask=True, use_drop_mask=False)
 
@@ -154,6 +156,55 @@ def test_deterministic():
         assert torch.equal(outs[0], other)
 
 
+def _compile_case(c_z, device):
+    """An inference-path TriMul closure plus an input factory, as the layer calls it."""
+    w = _weights(c_z, device)
+
+    def fn(pair, mask):
+        return esmfold2_trimul.triangle_multiplicative_update_with_residual(
+            pair, "outgoing", torch.zeros_like(pair), None, mask=mask, **w
+        )
+
+    def inputs(L):
+        pair = torch.randn(1, L, L, c_z, device=device).to(torch.bfloat16)
+        mask = torch.rand(1, L, L, device=device).to(torch.bfloat16)
+        return pair, mask
+
+    return fn, inputs
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_torch_compile_fullgraph_matches_eager():
+    """The layer declares `can_torch_compile`, so the op must trace without graph breaks."""
+    torch.manual_seed(0)
+    torch.compiler.reset()
+    fn, inputs = _compile_case(128, torch.device("cuda"))
+    pair, mask = inputs(24)
+
+    with torch.no_grad():
+        want = fn(pair, mask)
+        got = torch.compile(fn, fullgraph=True)(pair, mask)
+    torch.testing.assert_close(got, want)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_torch_compile_no_recompile_across_lengths():
+    """L must stay symbolic under dynamic shapes; protein lengths vary per call."""
+    torch.manual_seed(0)
+    torch.compiler.reset()
+    fn, inputs = _compile_case(128, torch.device("cuda"))
+    compiled = torch.compile(fn, fullgraph=True, dynamic=True)
+
+    with torch.no_grad():
+        compiled(*inputs(16))
+        with torch.compiler.set_stance("fail_on_recompile"):
+            for L in (24, 33, 64):
+                pair, mask = inputs(L)
+                torch.testing.assert_close(compiled(pair, mask), fn(pair, mask))
+
+
 @pytest.mark.kernels_ci
 @pytest.mark.parametrize("c_z", [16, 32, 96, 192])
 def test_rejects_unsupported_channels(c_z):
@@ -180,3 +231,5 @@ def test_layer_is_exposed_for_kernels():
     assert issubclass(layer, torch.nn.Module)
     # `kernels` requires layers to be stateless: no constructor of their own.
     assert layer.__init__ is torch.nn.Module.__init__
+    # Lets transformers register the layer for `Mode.TORCH_COMPILE`.
+    assert layer.can_torch_compile is True

esmfold2-trimul/torch-ext/esmfold2_trimul/fused_ln_residual.py

@@ -60,7 +60,9 @@ def _ln_fwd_kernel(
     TILE_M: tl.constexpr,
 ):
     pid = tl.program_id(axis=0).to(tl.int64)
-    M64 = M.to(tl.int64)
+    # tl.cast, not M.to: M is a plain int when Triton specializes it (M == 1),
+    # and when torch.compile analyzes the kernel.
+    M64 = tl.cast(M, tl.int64)
 
     offs_m = pid * TILE_M + tl.arange(0, TILE_M).to(tl.int64)
     offs_d = tl.arange(0, D).to(tl.int64)
@@ -108,7 +110,9 @@ def _ln_bwd_kernel(
     TILE_M: tl.constexpr,
 ):
     pid = tl.program_id(axis=0).to(tl.int64)
-    M64 = M.to(tl.int64)
+    # tl.cast, not M.to: M is a plain int when Triton specializes it (M == 1),
+    # and when torch.compile analyzes the kernel.
+    M64 = tl.cast(M, tl.int64)
 
     offs_m = pid * TILE_M + tl.arange(0, TILE_M).to(tl.int64)
     offs_d = tl.arange(0, D).to(tl.int64)

esmfold2-trimul/torch-ext/esmfold2_trimul/layers.py

@@ -33,6 +33,8 @@ def _bf16(t):
 
 
 class ESMFold2TriangleMultiplication(nn.Module):
+    can_torch_compile = True
+
     def forward(self, pair_grid: torch.Tensor, visibility: torch.Tensor | None = None) -> torch.Tensor:
         lat = self.dim
         pb = self.proj_bundle.weight

esmfold2-trimul/torch-ext/esmfold2_trimul/trimul_with_residual.py

@@ -76,8 +76,10 @@ def _gated_gemm_with_residual_kernel(
     M,
     N,
     K,
-    STRIDE_B: tl.constexpr,  # = N_ROW * N_COL  (for decoding b from m)
-    N_COL: tl.constexpr,  # = L_col  (for decoding j from m)
+    # Runtime args, not constexpr: both scale with sequence length, so constexpr
+    # would recompile the kernel (and torch.compile the graph) per distinct L.
+    STRIDE_B,  # = N_ROW * N_COL  (for decoding b from m)
+    N_COL,  # = L_col  (for decoding j from m)
     TILE_M: tl.constexpr,
     TILE_N: tl.constexpr,
     TILE_K: tl.constexpr,
@@ -215,8 +217,8 @@ def _gated_gemm_with_residual_backward_kernel(
     M,
     N,
     K,
-    STRIDE_B: tl.constexpr,
-    N_COL: tl.constexpr,
+    STRIDE_B,  # runtime, not constexpr: see the forward kernel
+    N_COL,
     TILE_M: tl.constexpr,
     TILE_N: tl.constexpr,
     TILE_K: tl.constexpr,
@@ -531,7 +533,6 @@ def _gated_gemm_with_residual(
     )
 
 
-@torch._dynamo.disable
 def triangle_multiplicative_update_with_residual(
     pair: torch.Tensor,
     direction: str,