replace triangle layers with nvidia cuequivariance kernels

bench_folding_grad.py

@@ -0,0 +1,120 @@
+"""Forward+backward benchmark: vanilla vs cueq-fast protenij.
+
+Uses an MSE-on-distogram-logits loss that exercises the full trunk
+(input embedder, MSA module, template embedder, Pairformer stack, distogram
+head). Diffusion + confidence head are excluded — they host no cueq-swapped
+kernels and the sampler adds stochasticity that complicates grad timing.
+
+This is the relevant regime for training: cueq's cuda-fwd/cuda-bwd kernels for
+triangle_attention (and triton VJP for triangle_multiplicative_update) get
+exercised end-to-end.
+"""
+import os
+os.environ["PROTENIX_DATA_ROOT_DIR"] = os.path.expanduser("~/.protenix")
+
+import argparse
+import statistics
+import time
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+
+from protenix.backend import load_model
+
+from bench_folding import SEQUENCES, enable_cueq, build_features
+
+
+def trunk_loss(model, features, key, n_cycle=4):
+    """MSE of distogram logits against a zero target. Scalar — grad OK.
+
+    recycle() stop_gradients between cycles, so only the last cycle contributes
+    to grad. We still run n_cycle forwards to match production wall time.
+    """
+    emb = model.embed_inputs(input_feature_dict=features)
+    trunk = model.recycle(
+        initial_embedding=emb, input_feature_dict=features,
+        recycling_steps=n_cycle, key=key,
+    )
+    logits = model.distogram_head(trunk.z)
+    return jnp.mean(jnp.square(logits))
+
+
+def _block(tree):
+    jax.tree.map(
+        lambda v: v.block_until_ready() if hasattr(v, "block_until_ready") else None,
+        tree,
+    )
+
+
+def _time(fn, *, n_iter):
+    # First call includes compile; time it separately.
+    t0 = time.perf_counter()
+    out = fn()
+    _block(out)
+    compile_plus_run = time.perf_counter() - t0
+
+    times = []
+    for _ in range(n_iter):
+        t0 = time.perf_counter()
+        out = fn()
+        _block(out)
+        times.append(time.perf_counter() - t0)
+    return compile_plus_run, times
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default="protenix_base_default_v1.0.0")
+    parser.add_argument("--n_cycle", type=int, default=4)
+    parser.add_argument("--n_iter", type=int, default=3)
+    parser.add_argument("--msa_dir", default="output_test_predict/msa",
+                        help="mmseqs2 MSA cache dir; set to '' for dummy MSA")
+    args = parser.parse_args()
+
+    print(f"device: {jax.devices()[0]}")
+    print(f"loading {args.model}...")
+    baseline = load_model(args.model)
+    cueq = enable_cueq(baseline)
+
+    msa_dir = args.msa_dir if args.msa_dir else None
+    n_cycle = args.n_cycle
+
+    def _fwd(model, features, key):
+        return trunk_loss(model, features, key, n_cycle=n_cycle)
+
+    fwd = eqx.filter_jit(_fwd)
+    fwd_bwd = eqx.filter_jit(eqx.filter_value_and_grad(_fwd))
+
+    print(f"\nn_cycle={n_cycle}  n_iter={args.n_iter}  "
+          f"msa={'real (mmseqs2)' if msa_dir else 'dummy'}")
+    print("loss = mean(distogram_logits ** 2)")
+
+    for name, seq in SEQUENCES.items():
+        features = build_features(seq, name, msa_dir=msa_dir)
+        n_tok = int(features["token_index"].shape[0])
+        print(f"\n=== {name} ({len(seq)} residues, {n_tok} tokens, "
+              f"msa={tuple(features['msa'].shape)}) ===")
+        key = jax.random.PRNGKey(0)
+
+        results = {}
+        for label, model in [("baseline", baseline), ("cueq-fast", cueq)]:
+            c_fwd, t_fwd = _time(lambda: fwd(model, features, key), n_iter=args.n_iter)
+            c_bb,  t_bb  = _time(lambda: fwd_bwd(model, features, key), n_iter=args.n_iter)
+            mf = statistics.mean(t_fwd) * 1e3
+            mb = statistics.mean(t_bb)  * 1e3
+            results[label] = (mf, mb, mb - mf)
+            print(f"  {label:10s}  "
+                  f"fwd mean={mf:6.1f}ms (compile+1st={c_fwd:6.2f}s)   "
+                  f"fwd+bwd mean={mb:6.1f}ms (compile+1st={c_bb:6.2f}s)   "
+                  f"bwd≈{mb - mf:6.1f}ms")
+
+        mf_b, mfb_b, bwd_b = results["baseline"]
+        mf_c, mfb_c, bwd_c = results["cueq-fast"]
+        print(f"  speedup (cueq vs baseline): "
+              f"fwd={mf_b/mf_c:.2f}x   bwd≈{bwd_b/bwd_c:.2f}x   "
+              f"fwd+bwd={mfb_b/mfb_c:.2f}x")
+
+
+if __name__ == "__main__":
+    main()

protenix/protenij.py

@@ -33,6 +33,25 @@
     import protenix.model.protenix
     import protenix.openfold_local.model.primitives
 
+
+# Monkey-patch cuequivariance_jax 0.8.x: its triangle_attention custom_vjp fwd
+# rule returns (a, lse, amax) as a tuple, but the primitive binding returns a
+# list — JAX's custom_vjp matcher rejects the mismatch during backward tracing.
+# Harmless if cueq isn't installed or is never used for a backward pass.
+try:
+    from cuequivariance_jax.triangle._triangle_attention import (
+        triangle_attention_custom_vjp as _cueq_ta_cvjp,
+    )
+    _cueq_orig_fwd = _cueq_ta_cvjp.fwd
+
+    def _cueq_patched_fwd(*args, **kwargs):
+        primal, residuals = _cueq_orig_fwd(*args, **kwargs)
+        return list(primal), residuals
+
+    _cueq_ta_cvjp.fwd = _cueq_patched_fwd
+except ImportError:
+    pass
+
 def move_final_dim_to_dim(x, dim: int):
     # permute_final_dims
     n_dim = len(x.shape)
@@ -283,21 +302,40 @@ class TriangleMultiplication(AbstractFromTorch):
     linear_z: backend.Linear
     sigmoid: any
     _outgoing: bool
+    # Static field (part of treedef) so stacking/partitioning never treats the
+    # bool as a dynamic leaf — older skeletons default it to False on load.
+    use_cueq: bool = eqx.field(default=False, static=True)
+
+    @staticmethod
+    def from_torch(m):
+        return TriangleMultiplication(
+            linear_a_p=from_torch(m.linear_a_p),
+            linear_a_g=from_torch(m.linear_a_g),
+            linear_b_p=from_torch(m.linear_b_p),
+            linear_b_g=from_torch(m.linear_b_g),
+            layer_norm_in=from_torch(m.layer_norm_in),
+            layer_norm_out=from_torch(m.layer_norm_out),
+            linear_g=from_torch(m.linear_g),
+            linear_z=from_torch(m.linear_z),
+            sigmoid=from_torch(m.sigmoid),
+            _outgoing=m._outgoing,
+        )
 
     def __call__(
         self,
         z_in: Float[Array, "... N N C_z"],
         mask: Bool[Array, "... N N"] | None,
     ) -> Float[Array, "... N N C_z"]:
-        #with jax.default_matmul_precision("float32"):
         if mask is None:
             mask = jnp.ones_like(z_in, shape=z_in.shape[:-1])
 
-        mask = mask[..., None]
+        if getattr(self, "use_cueq", False):
+            return self._cueq_forward(z_in, mask)
 
+        mask_e = mask[..., None]
         z = self.layer_norm_in(z_in)
-        a = self.linear_a_p(z) * mask * self.sigmoid(self.linear_a_g(z))
-        b = self.linear_b_p(z) * mask * self.sigmoid(self.linear_b_g(z))
+        a = self.linear_a_p(z) * mask_e * self.sigmoid(self.linear_a_g(z))
+        b = self.linear_b_p(z) * mask_e * self.sigmoid(self.linear_b_g(z))
 
         ### permute_final_dims (2, 0, 1)
         if self._outgoing:
@@ -307,7 +345,6 @@ def __call__(
             a = einops.rearrange(a, "... a b c -> ... c b a")
             b = einops.rearrange(b, "... a b c -> ... c a b")
         p = a @ b
-        # return p
         x = einops.rearrange(p, "... a b c -> ... b c a")
 
         x = self.layer_norm_out(x)
@@ -317,6 +354,47 @@ def __call__(
 
         return x * g
 
+    def _cueq_forward(self, z_in, mask):
+        """Triangle multiplicative update via NVIDIA cuequivariance JAX kernel.
+
+        Re-uses existing Linear/LayerNorm weights — cueq wants a/b projections
+        concatenated along the output axis (2*D_in, D_in). No new parameters.
+        """
+        import cuequivariance_jax as cuex
+
+        def stack_bias(ba, bb):
+            if ba is None and bb is None:
+                return None
+            if ba is None:
+                ba = jnp.zeros_like(bb)
+            if bb is None:
+                bb = jnp.zeros_like(ba)
+            return jnp.concatenate([ba, bb], axis=0)
+
+        p_in_w = jnp.concatenate([self.linear_a_p.weight, self.linear_b_p.weight], axis=0)
+        g_in_w = jnp.concatenate([self.linear_a_g.weight, self.linear_b_g.weight], axis=0)
+        p_in_b = stack_bias(self.linear_a_p.bias, self.linear_b_p.bias)
+        g_in_b = stack_bias(self.linear_a_g.bias, self.linear_b_g.bias)
+
+        return cuex.triangle_multiplicative_update(
+            z_in,
+            direction="outgoing" if self._outgoing else "incoming",
+            mask=mask,
+            norm_in_weight=self.layer_norm_in.weight,
+            norm_in_bias=self.layer_norm_in.bias,
+            p_in_weight=p_in_w,
+            p_in_bias=p_in_b,
+            g_in_weight=g_in_w,
+            g_in_bias=g_in_b,
+            norm_out_weight=self.layer_norm_out.weight,
+            norm_out_bias=self.layer_norm_out.bias,
+            p_out_weight=self.linear_z.weight,
+            p_out_bias=self.linear_z.bias,
+            g_out_weight=self.linear_g.weight,
+            g_out_bias=self.linear_g.bias,
+            eps=self.layer_norm_in.eps,
+        )
+
 
 class Attention(AbstractFromTorch):
     c_q: int  # input dimension of query
@@ -388,6 +466,18 @@ class TriangleAttention(AbstractFromTorch):
     linear: Linear
     mha: Attention
     inf: float
+    # Static: see TriangleMultiplication.use_cueq for rationale.
+    use_cueq: bool = eqx.field(default=False, static=True)
+
+    @staticmethod
+    def from_torch(m):
+        return TriangleAttention(
+            starting=m.starting,
+            layer_norm=from_torch(m.layer_norm),
+            linear=from_torch(m.linear),
+            mha=from_torch(m.mha),
+            inf=m.inf,
+        )
 
     def __call__(
         self,
@@ -403,21 +493,70 @@ def __call__(
 
         x = self.layer_norm(x)
 
-        # [*, I, 1, 1, J]
-        mask_bias = (self.inf * (mask - 1))[..., :, None, None, :]
-
         # [*, H, I, J]
         triangle_bias = einops.rearrange(self.linear(x), "... A B C -> ... C A B")
         # [*, 1, H, I, J]
         triangle_bias = einops.rearrange(triangle_bias, "... H I J -> ... 1 H I J")
 
-        x = self.mha(q_x=x, kv_x=x, biases=[mask_bias, triangle_bias])
+        if getattr(self, "use_cueq", False):
+            x = self._cueq_forward(x, mask, triangle_bias)
+        else:
+            # [*, I, 1, 1, J]
+            mask_bias = (self.inf * (mask - 1))[..., :, None, None, :]
+            x = self.mha(q_x=x, kv_x=x, biases=[mask_bias, triangle_bias])
 
         if not self.starting:
             x = einops.rearrange(x, "... I J C -> ... J I C")
 
         return x
 
+    def _cueq_forward(self, x, mask, triangle_bias):
+        """Triangle attention via NVIDIA cuequivariance JAX kernel.
+
+        Reuses mha's linear projections, gate, and output projection; only the
+        attention kernel itself is swapped. Checkpoint weights are unchanged.
+        """
+        import cuequivariance_jax as cuex
+
+        mha = self.mha
+        H = mha.no_heads
+        D = mha.c_hidden
+
+        q = einops.rearrange(mha.linear_q(x), "... N S (H D) -> ... N H S D", H=H)
+        k = einops.rearrange(mha.linear_k(x), "... N S (H D) -> ... N H S D", H=H)
+        v = einops.rearrange(mha.linear_v(x), "... N S (H D) -> ... N H S D", H=H)
+
+        # cueq expects rank-5 [B, N, H, S, D]; add a B=1 if absent.
+        added_batch = q.ndim == 4
+        if added_batch:
+            q, k, v = q[None], k[None], v[None]
+
+        bias = triangle_bias if triangle_bias.ndim == 5 else triangle_bias[None]
+        mask_cueq = mask.astype(bool)[..., :, None, None, :]
+        if mask_cueq.ndim == 4:
+            mask_cueq = mask_cueq[None]
+
+        scale = 1.0 / math.sqrt(D)
+        # TF32 / tensor-core path. HIGHEST (IEEE fp32) lacks a backward kernel
+        # in cueq 0.8.x and only matches baseline precision, which is itself
+        # TF32 on Ampere+, so DEFAULT is the only mode worth supporting.
+        out = cuex.triangle_attention(
+            q, k, v, bias, mask_cueq, scale, precision=jax.lax.Precision.DEFAULT,
+        )
+        # cueq returns [output, logsumexp, max_val]; keep just the output.
+        o = out[0] if isinstance(out, (tuple, list)) else out
+
+        if added_batch:
+            o = o[0]
+
+        o = einops.rearrange(o, "... N H S D -> ... N S H D")
+        if mha.linear_g is not None:
+            g = jax.nn.sigmoid(mha.linear_g(x))
+            g = einops.rearrange(g, "... N S (H D) -> ... N S H D", H=H)
+            o = o * g
+        o = einops.rearrange(o, "... N S H D -> ... N S (H D)")
+        return mha.linear_o(o)
+
 
 def _attention(
     q,

pyproject.toml

@@ -10,7 +10,7 @@ dependencies = [
     "einops>=0.8.0",
     "equinox>=0.13.0",
     "gemmi>=0.6.5",
-    "jax>=0.6.2",
+    "jax>=0.8.1",
     "ml-collections>=1.1.0",
     "modelcif>=0.7",
     "numpy>=2.1.0",
@@ -26,6 +26,9 @@ dependencies = [
     "setuptools>=80.9.0",
     "tqdm>=4.67.1",
     "huggingface_hub>=0.20.0",
+    "cuequivariance-jax>=0.8.1,<0.9",
+    "cuequivariance-ops-jax-cu12>=0.8.1,<0.9",
+    "jax-triton>=0.3.1",
 ]
 
 [project.optional-dependencies]

test_predict.py

@@ -48,7 +48,7 @@
 
 eqx_path = os.path.join(CACHE_DIR, f"{MODEL_NAME}")
 print(f"Loading JAX model from {eqx_path}...")
-jax_model = load_model(eqx_path)
+jax_model = load_model(MODEL_NAME)
 
 # Override diffusion parameters for vanilla ODE sampling
 jax_model = eqx.tree_at(lambda m: m.gamma0, jax_model, 0.0)