From 9eda5d5f93c4b26971e961dabb82d7b2ec5c8837 Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 5 Jun 2026 19:23:12 +0000 Subject: [PATCH 1/5] Document triangle attention kv_lengths fast path --- CHANGELOG.md | 5 +- .../cuequivariance_torch/SKILL.md | 5 +- .../cuequivariance_torch/__init__.py | 2 + .../primitives/triangle.py | 76 ++++++++++++++++--- .../tests/primitives/triangle_test.py | 50 ++++++++++++ docs/api/cuequivariance_torch.rst | 1 + 6 files changed, 128 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 31f45f5..d08e47e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,9 @@ ## Latest Changes ## (unreleased) +### Added +- [Torch] `cuet.triangle_attention` accepts `kv_lengths`, an int32 per-row key/value length tensor for right-padded sequence masks. Passing `kv_lengths` selects the Blackwell sm100f length fast path when available; dense `mask` remains the correct fallback path for arbitrary masks. Added `cuet.mask_to_kv_lengths` to convert and validate prefix masks. + ## 0.10.0 (2026-04-21) ### Added @@ -282,4 +285,4 @@ output = cuex.equivariant_tensor_product(e, w, input) ## 0.1.0 (2024-11-18) -- Beta version of cuEquivariance released. \ No newline at end of file +- Beta version of cuEquivariance released. diff --git a/cuequivariance_torch/cuequivariance_torch/SKILL.md b/cuequivariance_torch/cuequivariance_torch/SKILL.md index e2a92e4..499e5d2 100644 --- a/cuequivariance_torch/cuequivariance_torch/SKILL.md +++ b/cuequivariance_torch/cuequivariance_torch/SKILL.md @@ -306,7 +306,10 @@ Require `cuequivariance_ops_torch`. ```python # Triangle attention with pair bias -out = cuet.triangle_attention(q, k, v, bias, mask=mask, scale=scale) +kv_lengths = cuet.mask_to_kv_lengths(prefix_mask) +out = cuet.triangle_attention(q, k, v, bias, scale=scale, kv_lengths=kv_lengths) +# Arbitrary dense masks are supported through the fallback path: +out = cuet.triangle_attention(q, k, v, bias, mask=holey_mask, scale=scale) # q, k, v: (B, N, H, Q/K, D), bias: (B, 1, H, Q, K) # Triangle multiplicative update diff --git a/cuequivariance_torch/cuequivariance_torch/__init__.py b/cuequivariance_torch/cuequivariance_torch/__init__.py index 44fcd85..56293eb 100644 --- a/cuequivariance_torch/cuequivariance_torch/__init__.py +++ b/cuequivariance_torch/cuequivariance_torch/__init__.py @@ -42,6 +42,7 @@ ) from .operations.spherical_harmonics import SphericalHarmonics from .primitives.triangle import ( + mask_to_kv_lengths, triangle_attention, triangle_multiplicative_update, attention_pair_bias, @@ -96,6 +97,7 @@ def register_tensorrt_plugins(): "vector_to_euler_angles", "Inversion", "SphericalHarmonics", + "mask_to_kv_lengths", "triangle_attention", "triangle_multiplicative_update", "attention_pair_bias", diff --git a/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py b/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py index 70c72c2..b425ebd 100644 --- a/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py +++ b/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py @@ -39,6 +39,44 @@ class TriMulPrecision(enum.IntEnum): # type: ignore BFLOAT16 = 5 +def mask_to_kv_lengths(mask: torch.Tensor) -> torch.Tensor: + r""" + Convert a right-padded triangle-attention mask to per-row key/value lengths. + + Args: + mask (torch.Tensor): Boolean-like mask of shape (B, N, 1, 1, K). For B=1, + can also be (N, 1, 1, K). The last dimension must be prefix-shaped: + all True entries first, followed by all False entries. + + Returns: + torch.Tensor: int32 tensor of shape (B, N, 1, 1, 1) containing each row's + effective key/value length. Zero-length rows are allowed. + + Raises: + ValueError: If the mask shape is not supported or the mask is not prefix-shaped. + + Note: + This helper validates the prefix contract. For torch.compile-heavy code, compute + lengths before the compiled region and pass them to :func:`triangle_attention` + with ``kv_lengths=...``. + """ + mask_bool = mask.to(dtype=torch.bool) + while len(mask_bool.shape) < 5: + mask_bool = mask_bool.unsqueeze(0) + if mask_bool.ndim != 5 or mask_bool.shape[2:4] != (1, 1): + raise ValueError( + "mask_to_kv_lengths: mask must have shape (B, N, 1, 1, K) " + f"after adding leading singleton dimensions, got {tuple(mask_bool.shape)}" + ) + lengths = mask_bool.to(dtype=torch.int32).sum(dim=-1, keepdim=True).to(torch.int32) + prefix = torch.arange(mask_bool.shape[-1], device=mask_bool.device).view( + 1, 1, 1, 1, mask_bool.shape[-1] + ) < lengths + if not bool(torch.all(mask_bool == prefix).item()): + raise ValueError("mask_to_kv_lengths: mask must be right-padded/prefix-shaped") + return lengths.detach().contiguous() + + def triangle_attention( q: torch.Tensor, k: torch.Tensor, @@ -47,6 +85,8 @@ def triangle_attention( mask: Optional[torch.Tensor] = None, scale: Optional[float] = None, return_aux: bool = False, + *, + kv_lengths: Optional[torch.Tensor] = None, ) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: r""" Triangle Attention @@ -68,10 +108,17 @@ def triangle_attention( capability 10.0 or 10.3), will be cast to match q/k/v dtype (bf16/fp16) for best performance. mask (torch.Tensor, optional): Mask tensor of shape (B, N, 1, 1, K). For B=1, can also be (N, 1, 1, K). - Will be cast to bool internally. + Will be cast to bool internally. Dense masks accept any pattern and route to the + correct fallback path. For right-padded masks, pass ``kv_lengths`` instead to use + the Blackwell sm100f length fast path. scale (float, optional): Float scale for q (s in the equation). If None, value 1/sqrt(d) is used. return_aux (bool): If True, two auxiliary tensors are returned along with the result. Defaults to False. + kv_lengths (torch.Tensor, optional): int32 tensor of shape (B, N, 1, 1, 1) + containing each row's effective K length. This represents a right-padded / + prefix mask: positions ``j < kv_lengths[b, n]`` are valid and later positions + are masked. Pass either ``mask`` or ``kv_lengths``, not both. On supported + Blackwell cu13 builds, ``kv_lengths`` selects the sm100f length fast path. Note: - B: batch size @@ -90,12 +137,12 @@ def triangle_attention( (1) Context is saved for backward pass. You don't need to save it manually. (2) Kernel precision (fp32, bf16, fp16) is based on input dtypes. For tf32, set it from torch global scope (3) Triangle attention kernel supports: all hidden_dim<=32 and divisible by 4 for tf32/fp32, and for all hidden_dim<=128 and divisible by 8 for bf16/fp16 (standard kernels). On Blackwell GPUs (compute capability 10.0 or 10.3), the sm100f kernel supports hidden_dim<=256 for forward passes and hidden_dim<=128 for backward passes. In the rare instance that the kernel does not support an input config, fallback to torch is enabled instead of erroring out. - (4) Blackwell-optimized kernels (for compute capabilities 10.0 and 10.3) provide superior performance especially for long sequences and higher head dimensions. These kernels require the sequence length N to be a multiple of 8 for the forward pass; pad the sequence if necessary. The kernel provides optimal performance for a "padding mask" consisting in (all True, followed by all False) in the last dimension. Currently, this feature is supported only for cu13 builds. + (4) Blackwell-optimized kernels (for compute capabilities 10.0 and 10.3) provide superior performance especially for long sequences and higher head dimensions. These kernels require the key/value sequence length K to be a multiple of 8 for the forward pass; pad the sequence if necessary. Use ``kv_lengths`` for right-padded sequence masks to select the sm100f length fast path. A dense ``mask`` without ``kv_lengths`` remains correct for arbitrary or holey patterns, but it routes to the fallback path. Example: >>> import torch >>> import math - >>> from cuequivariance_torch import triangle_attention + >>> from cuequivariance_torch import mask_to_kv_lengths, triangle_attention >>> if torch.cuda.is_available(): # doctest: +SKIP ... device = torch.device("cuda") ... # Set up dimensions @@ -109,14 +156,25 @@ def triangle_attention( ... device=device, dtype=torch.float16, requires_grad=True) ... bias = torch.randn(batch_size, 1, num_heads, seq_len, seq_len, ... device=device, dtype=torch.float32, requires_grad=True) - ... # Create optional mask - ... mask = torch.rand(batch_size, seq_len, 1, 1, seq_len, - ... device=device) < 0.5 + ... # Right-padded sequence mask: valid tokens first, padding last. + ... seq_lengths = torch.tensor([96], device=device, dtype=torch.int32) + ... positions = torch.arange(seq_len, device=device) + ... row_valid = positions.view(1, seq_len) < seq_lengths.view(batch_size, 1) + ... col_valid = positions.view(1, seq_len) < seq_lengths.view(batch_size, 1) + ... mask = row_valid.view(batch_size, seq_len, 1, 1, 1) & col_valid.view( + ... batch_size, 1, 1, 1, seq_len) + ... kv_lengths = mask_to_kv_lengths(mask) ... # Calculate scale ... scale = 1 / math.sqrt(hidden_dim) - ... # Forward pass + ... # Forward pass using the Blackwell sm100f length fast path when available. ... output, lse, max_val = triangle_attention( - ... q=q, k=k, v=v, bias=bias, mask=mask, scale=scale, return_aux=True) + ... q=q, k=k, v=v, bias=bias, scale=scale, return_aux=True, + ... kv_lengths=kv_lengths) + ... # Arbitrary dense masks are still correct; they use the fallback path. + ... arbitrary_mask = torch.rand(batch_size, seq_len, 1, 1, seq_len, + ... device=device) < 0.5 + ... fallback_output = triangle_attention( + ... q=q, k=k, v=v, bias=bias, mask=arbitrary_mask, scale=scale) ... print(output.shape) # torch.Size([1, 128, 2, 128, 32]) ... # Create gradient tensor and perform backward pass ... grad_out = torch.randn_like(output) @@ -140,7 +198,7 @@ def triangle_attention( "Error importing triangle_attention from cuequivariance_ops_torch." ) else: - return f(q, k, v, bias, mask, scale, return_aux) + return f(q, k, v, bias, mask, scale, return_aux, kv_lengths=kv_lengths) def triangle_multiplicative_update( diff --git a/cuequivariance_torch/tests/primitives/triangle_test.py b/cuequivariance_torch/tests/primitives/triangle_test.py index 4f0a6eb..a82c716 100644 --- a/cuequivariance_torch/tests/primitives/triangle_test.py +++ b/cuequivariance_torch/tests/primitives/triangle_test.py @@ -94,6 +94,56 @@ def test_triangle_attention(): ) +def test_triangle_attention_kv_lengths(): + if torch.cuda.is_available(): + device = torch.device("cuda") + batch_size, seq_len, num_heads, hidden_dim = 1, 16, 2, 32 + q = torch.randn( + batch_size, + seq_len, + num_heads, + seq_len, + hidden_dim, + device=device, + dtype=torch.float16, + ) + k = torch.randn_like(q) + v = torch.randn_like(q) + bias = torch.randn( + batch_size, + 1, + num_heads, + seq_len, + seq_len, + device=device, + dtype=torch.float32, + ) + seq_lengths = torch.tensor([12], device=device, dtype=torch.int32) + positions = torch.arange(seq_len, device=device) + row_valid = positions.view(1, seq_len) < seq_lengths.view(batch_size, 1) + col_valid = positions.view(1, seq_len) < seq_lengths.view(batch_size, 1) + mask = row_valid.view(batch_size, seq_len, 1, 1, 1) & col_valid.view( + batch_size, 1, 1, 1, seq_len + ) + kv_lengths = cuet.mask_to_kv_lengths(mask) + + assert kv_lengths.shape == torch.Size([batch_size, seq_len, 1, 1, 1]) + assert kv_lengths.dtype == torch.int32 + output = cuet.triangle_attention( + q=q, + k=k, + v=v, + bias=bias, + scale=1 / math.sqrt(hidden_dim), + kv_lengths=kv_lengths, + ) + assert output.shape == torch.Size( + [batch_size, seq_len, num_heads, seq_len, hidden_dim] + ) + zero_rows = kv_lengths.view(batch_size, seq_len) == 0 + torch.testing.assert_close(output[zero_rows], torch.zeros_like(output[zero_rows])) + + def test_triangle_multiplicative_update(): if torch.cuda.is_available(): device = torch.device("cuda") diff --git a/docs/api/cuequivariance_torch.rst b/docs/api/cuequivariance_torch.rst index 839bcbb..7e0a47c 100644 --- a/docs/api/cuequivariance_torch.rst +++ b/docs/api/cuequivariance_torch.rst @@ -68,6 +68,7 @@ Triangle :toctree: generated/ :template: function_template.rst + mask_to_kv_lengths triangle_attention triangle_multiplicative_update attention_pair_bias From 4c81937c906c6df67c42fe122fdedcf4d53e7377 Mon Sep 17 00:00:00 2001 From: Jonathan Mitchell Date: Fri, 5 Jun 2026 14:31:24 -0700 Subject: [PATCH 2/5] 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 Co-Authored-By: Claude Opus 4.8 --- .../cuequivariance_torch/primitives/triangle.py | 7 +++++++ .../tests/primitives/triangle_test.py | 16 ++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py b/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py index b425ebd..7b74a42 100644 --- a/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py +++ b/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: diff --git a/cuequivariance_torch/tests/primitives/triangle_test.py b/cuequivariance_torch/tests/primitives/triangle_test.py index a82c716..8917f2f 100644 --- a/cuequivariance_torch/tests/primitives/triangle_test.py +++ b/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") From 8fb03699cad754c8bd332ed6cdef31c9eab6bd18 Mon Sep 17 00:00:00 2001 From: Jonathan Mitchell Date: Thu, 11 Jun 2026 14:12:41 -0700 Subject: [PATCH 3/5] Document the kv_lengths <= K clamp contract in triangle_attention State that kv_lengths must be in [0, K]; values greater than K are clamped to K (the row then attends the full key sequence), matching the backend clamp that prevents an out-of-bounds read on the SM100f length fast path. Signed-off-by: Jonathan Mitchell Co-Authored-By: Claude Opus 4.8 --- .../cuequivariance_torch/primitives/triangle.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py b/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py index 7b74a42..f8c4000 100644 --- a/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py +++ b/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py @@ -119,6 +119,8 @@ def triangle_attention( prefix mask: positions ``j < kv_lengths[b, n]`` are valid and later positions are masked. Pass either ``mask`` or ``kv_lengths``, not both. On supported Blackwell cu13 builds, ``kv_lengths`` selects the sm100f length fast path. + Each length must be in ``[0, K]``; values greater than ``K`` are clamped + to ``K`` (the row then attends the full key sequence). Note: - B: batch size From e07c160cb503dfc2785f63506c0cacbe7c7e47a9 Mon Sep 17 00:00:00 2001 From: Jonathan Mitchell Date: Thu, 11 Jun 2026 14:36:32 -0700 Subject: [PATCH 4/5] Apply ruff-format to mask_to_kv_lengths (style check) ruff-format (v0.13.1, the repo's pinned pre-commit version) parenthesizes the prefix comparison; the committed form failed the format check. No behavior change. Signed-off-by: Jonathan Mitchell Co-Authored-By: Claude Opus 4.8 --- .../cuequivariance_torch/primitives/triangle.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py b/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py index f8c4000..3affc14 100644 --- a/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py +++ b/cuequivariance_torch/cuequivariance_torch/primitives/triangle.py @@ -69,9 +69,12 @@ def mask_to_kv_lengths(mask: torch.Tensor) -> torch.Tensor: f"after adding leading singleton dimensions, got {tuple(mask_bool.shape)}" ) lengths = mask_bool.to(dtype=torch.int32).sum(dim=-1, keepdim=True).to(torch.int32) - prefix = torch.arange(mask_bool.shape[-1], device=mask_bool.device).view( - 1, 1, 1, 1, mask_bool.shape[-1] - ) < lengths + prefix = ( + torch.arange(mask_bool.shape[-1], device=mask_bool.device).view( + 1, 1, 1, 1, mask_bool.shape[-1] + ) + < lengths + ) if not bool(torch.all(mask_bool == prefix).item()): raise ValueError("mask_to_kv_lengths: mask must be right-padded/prefix-shaped") return lengths.detach().contiguous() From 0aa1562687ddf1c19e4d46390c127cfeaeb7875c Mon Sep 17 00:00:00 2001 From: Jonathan Mitchell Date: Thu, 11 Jun 2026 14:38:34 -0700 Subject: [PATCH 5/5] Apply ruff-format to __init__.py and triangle_test.py (style check) ruff-format (v0.13.1, the repo's pinned pre-commit version): a blank line after the module docstring in __init__.py, and wrapping one long assert_close in triangle_test.py. No behavior change. With this, all #287-changed Python files pass ruff check + format. Signed-off-by: Jonathan Mitchell Co-Authored-By: Claude Opus 4.8 --- cuequivariance_torch/cuequivariance_torch/__init__.py | 1 + cuequivariance_torch/tests/primitives/triangle_test.py | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/cuequivariance_torch/cuequivariance_torch/__init__.py b/cuequivariance_torch/cuequivariance_torch/__init__.py index 56293eb..39a15de 100644 --- a/cuequivariance_torch/cuequivariance_torch/__init__.py +++ b/cuequivariance_torch/cuequivariance_torch/__init__.py @@ -16,6 +16,7 @@ For AI coding assistants: run `python -m cuequivariance_torch skill` for detailed usage guidance. """ + import importlib.resources __version__ = ( diff --git a/cuequivariance_torch/tests/primitives/triangle_test.py b/cuequivariance_torch/tests/primitives/triangle_test.py index 8917f2f..b0207f5 100644 --- a/cuequivariance_torch/tests/primitives/triangle_test.py +++ b/cuequivariance_torch/tests/primitives/triangle_test.py @@ -142,7 +142,9 @@ def test_triangle_attention_kv_lengths(): [batch_size, seq_len, num_heads, seq_len, hidden_dim] ) zero_rows = kv_lengths.view(batch_size, seq_len) == 0 - torch.testing.assert_close(output[zero_rows], torch.zeros_like(output[zero_rows])) + torch.testing.assert_close( + output[zero_rows], torch.zeros_like(output[zero_rows]) + ) def test_triangle_attention_mask_and_kv_lengths_mutually_exclusive():