Guard triangle_attention against passing mask and kv_lengths together

Raise ValueError in the public wrapper when both `mask` and `kv_lengths` are
supplied. They are mutually exclusive: `kv_lengths` selects the SM100f length
fast path, while a dense `mask` uses the fallback path. The guard sits before
the backend dispatch, so it surfaces the error early and needs no GPU. Add a
GPU-free regression test.

Signed-off-by: Jonathan Mitchell <jomitchell@nvidia.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: jomitchellnv

cuequivariance_torch/cuequivariance_torch/primitives/triangle.py

@@ -191,6 +191,13 @@ def triangle_attention(
         torch.Size([1, 1, 2, 128, 128])
     """
 
+    if mask is not None and kv_lengths is not None:
+        raise ValueError(
+            "triangle_attention: pass either `mask` or `kv_lengths`, not both. "
+            "`kv_lengths` selects the SM100f length fast path; a dense `mask` uses "
+            "the fallback path."
+        )
+
     try:
         from cuequivariance_ops_torch import triangle_attention as f
     except Exception:

cuequivariance_torch/tests/primitives/triangle_test.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 import math
 
+import pytest
 import torch
 
 import cuequivariance_torch as cuet
@@ -144,6 +145,21 @@ def test_triangle_attention_kv_lengths():
         torch.testing.assert_close(output[zero_rows], torch.zeros_like(output[zero_rows]))
 
 
+def test_triangle_attention_mask_and_kv_lengths_mutually_exclusive():
+    # kv_lengths selects the SM100f length fast path; a dense mask uses the
+    # fallback. Passing both is contradictory and must raise. The guard lives in
+    # the public wrapper, before any backend dispatch, so this needs no GPU.
+    n = 8
+    q = torch.zeros(1, n, 1, n, 8)
+    k = torch.zeros(1, n, 1, n, 8)
+    v = torch.zeros(1, n, 1, n, 8)
+    bias = torch.zeros(1, 1, 1, n, n)
+    mask = torch.ones(1, n, 1, 1, n, dtype=torch.bool)
+    kv_lengths = torch.full((1, n, 1, 1, 1), n, dtype=torch.int32)
+    with pytest.raises(ValueError, match="not both"):
+        cuet.triangle_attention(q, k, v, bias, mask=mask, kv_lengths=kv_lengths)
+
+
 def test_triangle_multiplicative_update():
     if torch.cuda.is_available():
         device = torch.device("cuda")