Document triangle attention kv_lengths fast path

Author: codex

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.
+- Beta version of cuEquivariance released.

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

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",

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(

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")

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