Expose APB normalization kind in the torch wrapper

cuequivariance_torch/cuequivariance_torch/SKILL.md

@@ -327,8 +327,9 @@ out, _ = cuet.attention_pair_bias(
     w_proj_g, w_proj_o, w_proj_z,
     w_ln_z=w_ln_z, b_ln_z=b_ln_z,
 )
-# Optional generalized (Proteina/Complexa) params, all keyword-only and None by
-# default: b_proj_k, b_proj_v, w_ln_q, b_ln_q, w_ln_k, b_ln_k, b_proj_g.
+# Optional generalized projection params are keyword-only:
+# b_proj_k, b_proj_v, w_ln_q, b_ln_q, w_ln_k, b_ln_k, b_proj_g (default None),
+# and norm_kind (default "layer_norm").
 ```
 
 ## ONNX and TensorRT export

cuequivariance_torch/cuequivariance_torch/primitives/triangle.py

@@ -363,14 +363,13 @@ def attention_pair_bias(
     w_ln_k: Optional[torch.Tensor] = None,
     b_ln_k: Optional[torch.Tensor] = None,
     b_proj_g: Optional[torch.Tensor] = None,
+    norm_kind: str = "layer_norm",
 ):
     """Compute attention with a pairwise bias.
 
-    Takes the single representation ``single_repr`` and computes LayerNorm, the
-    Q/K/V and gating projections, the pair-bias term
-    ``Linear(LayerNorm(pair_repr))``, biased and masked softmax attention,
-    the output gate ``sigmoid(Linear(LayerNorm(single_repr)))``, and the output
-    projection.
+    Normalizes the single and pair representations according to ``norm_kind``,
+    then computes the Q/K/V and gating projections, pair-bias term, biased and
+    masked softmax attention, output gate, and output projection.
     The backend selects a supported, profitable optimized route for the current
     inputs and otherwise uses the native PyTorch implementation (see Notes). Q, K,
     V and the gate are all projected from the same normalized ``single_repr``
@@ -381,17 +380,17 @@ def attention_pair_bias(
     attention inner dim is ``H * DH``), ``z_dim`` pair feature dim.
 
     Args:
-        single_repr: Single/token representation of shape (B * M, N, D). LayerNorm
-            and the Q/K/V/gate projections are applied to this tensor inside the op
-            (the gate is computed from the normalized representation).
+        single_repr: Single/token representation of shape (B * M, N, D).
+            Normalization and the Q/K/V/gate projections are applied to this tensor
+            inside the op (the gate is computed from the normalized representation).
         pair_repr: Pairwise tensor of shape (B, N, N, z_dim). When
             ``is_cached_z_proj`` is True, ``pair_repr`` is instead the
             already-projected bias of shape (B, H, N, N).
         mask: Attention mask of shape (B, N) or (B * M, N) (0 = masked, 1 = valid).
             If None, all positions are treated as valid.
         num_heads: Number of attention heads.
-        w_ln_a: Weight for the LayerNorm of ``single_repr`` of shape (D,).
-        b_ln_a: Bias for the LayerNorm of ``single_repr`` of shape (D,). May be None.
+        w_ln_a: Normalization weight for ``single_repr`` of shape (D,).
+        b_ln_a: Normalization bias for ``single_repr`` of shape (D,). May be None.
         w_proj_q: Weight for the query projection of shape (H * DH, D).
         b_proj_q: Bias for the query projection of shape (H * DH,). May be None.
         w_proj_k: Weight for the key projection of shape (H * DH, D).
@@ -403,10 +402,10 @@ def attention_pair_bias(
             projected. Defaults to None.
         b_proj_o: Bias for the output projection of shape (D,). Defaults to None.
         b_proj_z: Bias for the pair projection of shape (H,). Defaults to None.
-        w_ln_z: Weight for the LayerNorm of ``pair_repr`` of shape (z_dim,). May be None.
-        b_ln_z: Bias for the LayerNorm of ``pair_repr`` of shape (z_dim,). May be None.
+        w_ln_z: Normalization weight for ``pair_repr`` of shape (z_dim,). May be None.
+        b_ln_z: Normalization bias for ``pair_repr`` of shape (z_dim,). May be None.
         inf: Large value used for masking invalid attention positions. Defaults to 1e6.
-        eps: Epsilon value for layer normalization. Defaults to 1e-5.
+        eps: Epsilon value for normalization. Defaults to 1e-5.
         attn_scale: Scaling factor for attention scores. If None, uses
             1/sqrt(head_dim). Defaults to None.
         return_z_proj: Whether to return the projected pair tensor as the second
@@ -424,6 +423,9 @@ def attention_pair_bias(
         b_ln_k: Optional bias for projected-key LayerNorm of shape (H * DH,).
             Weight and bias are independently optional. Defaults to None.
         b_proj_g: Optional gating-projection bias of shape (H * DH,). Defaults to None.
+        norm_kind: Normalization applied to the single, optional Q/K, and pair
+            representations. ``"layer_norm"`` preserves the existing behavior;
+            ``"rms_norm"`` requires every normalization bias to be ``None``.
 
     Returns:
         - output (torch.Tensor): Attention output of shape (B * M, N, D) with the
@@ -441,13 +443,17 @@ def attention_pair_bias(
           replicas can share one pair representation in a single forward pass.
     """
 
+    if w_proj_z is None and not is_cached_z_proj:
+        raise ValueError("w_proj_z is required unless is_cached_z_proj=True")
+
     try:
         from cuequivariance_ops_torch import attention_pair_bias as f
-    except Exception:
+    except Exception as error:
         raise ImportError(
             "Error importing attention_pair_bias from cuequivariance_ops_torch."
-        )
+        ) from error
     else:
+        norm_kind_kwargs = {} if norm_kind == "layer_norm" else {"norm_kind": norm_kind}
         return f(
             single_repr,
             pair_repr,
@@ -478,4 +484,5 @@ def attention_pair_bias(
             w_ln_k=w_ln_k,
             b_ln_k=b_ln_k,
             b_proj_g=b_proj_g,
+            **norm_kind_kwargs,
         )

cuequivariance_torch/tests/primitives/triangle_test.py

@@ -236,7 +236,7 @@ def test_attention_pair_bias():
 
 
 def test_attention_pair_bias_generalized_projection():
-    # The optional generalized (Proteina/Complexa) projection params must be
+    # The optional generalized projection params must be
     # forwarded to the backend and actually change the output relative to the
     # strict OpenFold3/Boltz call that omits them.
     if torch.cuda.is_available():
@@ -370,13 +370,35 @@ def backend_spy(
     for name, value in generalized.items():
         assert calls[-1]["kwargs"][name] is value
 
+    rms_generalized = {
+        **generalized,
+        "b_ln_q": None,
+        "b_ln_k": None,
+    }
     actual, _ = cuet.attention_pair_bias(
         **common,
         pair_repr=torch.tensor([11.0]),
+        w_proj_z=torch.tensor([12.0]),
+        norm_kind="rms_norm",
+        **rms_generalized,
+    )
+    assert actual is output
+    assert calls[-1]["kwargs"]["norm_kind"] == "rms_norm"
+
+    actual, _ = cuet.attention_pair_bias(
+        **common,
+        pair_repr=torch.tensor([13.0]),
         is_cached_z_proj=True,
     )
     assert actual is cached_output
     assert calls[-1]["w_proj_z"] is None
     assert calls[-1]["kwargs"]["is_cached_z_proj"] is True
+    assert "norm_kind" not in calls[-1]["kwargs"]
     for name in generalized:
         assert calls[-1]["kwargs"][name] is None
+
+    with pytest.raises(ValueError, match="w_proj_z is required"):
+        cuet.attention_pair_bias(
+            **common,
+            pair_repr=torch.tensor([14.0]),
+        )