diff --git a/esmfold2-trimul/README.md b/esmfold2-trimul/README.md index 2dee38532..bbe86eb61 100644 --- a/esmfold2-trimul/README.md +++ b/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 diff --git a/esmfold2-trimul/build.toml b/esmfold2-trimul/build.toml index ab82261e3..92d2a6a73 100644 --- a/esmfold2-trimul/build.toml +++ b/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] diff --git a/esmfold2-trimul/tests/test_esmfold2_trimul.py b/esmfold2-trimul/tests/test_esmfold2_trimul.py index e726374cb..bcc4ef2b2 100644 --- a/esmfold2-trimul/tests/test_esmfold2_trimul.py +++ b/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 diff --git a/esmfold2-trimul/torch-ext/esmfold2_trimul/fused_ln_residual.py b/esmfold2-trimul/torch-ext/esmfold2_trimul/fused_ln_residual.py index ec349e64a..26d9dfa32 100644 --- a/esmfold2-trimul/torch-ext/esmfold2_trimul/fused_ln_residual.py +++ b/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) diff --git a/esmfold2-trimul/torch-ext/esmfold2_trimul/layers.py b/esmfold2-trimul/torch-ext/esmfold2_trimul/layers.py index 246d0ecb4..cebc3d78e 100644 --- a/esmfold2-trimul/torch-ext/esmfold2_trimul/layers.py +++ b/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 diff --git a/esmfold2-trimul/torch-ext/esmfold2_trimul/trimul_with_residual.py b/esmfold2-trimul/torch-ext/esmfold2_trimul/trimul_with_residual.py index 88131eb68..2b5b112b5 100644 --- a/esmfold2-trimul/torch-ext/esmfold2_trimul/trimul_with_residual.py +++ b/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,