triangle_attention#

cuequivariance_torch.triangle_attention(
q,
k,
v,
bias,
mask=None,
scale=None,
return_aux=False,
*,
kv_lengths=None,
)#

Triangle Attention

\[\begin{split}\text{Attention}_q(Q, K, V, B, M) = \sum_k\left[\text{softmax}_k\left(\begin{cases} s\, Q_q \cdot K_k + B_{qk} & \text{if } M_k = 1 \\ -10^9 & \text{otherwise} \end{cases}\right) V_k \right]\end{split}\]
Parameters:

  • q (torch.Tensor) – Query tensor of shape (B, N, H, Q, D). For B=1, can also be (N, H, Q, D).

  • k (torch.Tensor) – Key tensor of shape (B, N, H, K, D). For B=1, can also be (N, H, K, D).

  • v (torch.Tensor) – Value tensor of shape (B, N, H, K, D). For B=1, can also be (N, H, K, D).

  • bias (torch.Tensor) – Bias tensor of shape (B, 1, H, Q, K), For B=1, can also be (1, H, Q, K). Will be cast to float32 for standard kernels. On Blackwell GPUs (sm100f, compute 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. 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. Each length must be in [0, K]; values greater than K are clamped to K (the row then attends the full key sequence).

Return type:

Tensor | Tuple[Tensor, Tensor, Tensor]

Note

  • B: batch size

  • N: number of tokens

  • H: number of heads

  • Q: number of query tokens

  • K: number of key tokens

  • D: attention dimension

Returns:

Output tensor of shape (B, N, H, Q, D). dtype=q.dtype - lse(torch.Tensor): Auxiliary result (for special use only). dtype=float32 - max(torch.Tensor): Auxiliary result (for special use only). dtype=float32

Return type:

  • output(torch.Tensor)

Parameters:

Notes

  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 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 mask_to_kv_lengths, triangle_attention
>>> if torch.cuda.is_available():
...     device = torch.device("cuda")
...     # Set up dimensions
...     batch_size, seq_len, num_heads, hidden_dim = 1, 128, 2, 32
...     # Create input tensors on GPU with float16 precision
...     q = torch.randn(batch_size, seq_len, num_heads, seq_len, hidden_dim,
...                     device=device, dtype=torch.float16, requires_grad=True)
...     k = torch.randn(batch_size, seq_len, num_heads, seq_len, hidden_dim,
...                     device=device, dtype=torch.float16, requires_grad=True)
...     v = torch.randn(batch_size, seq_len, num_heads, seq_len, hidden_dim,
...                     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)
...     # 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 using the Blackwell sm100f length fast path when available.
...     output, lse, max_val = triangle_attention(
...         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)
...     output.backward(grad_out)
...     # Access gradients
...     print(q.grad.shape)  # torch.Size([1, 128, 2, 128, 32])
...     print(k.grad.shape)  # torch.Size([1, 128, 2, 128, 32])
...     print(v.grad.shape)  # torch.Size([1, 128, 2, 128, 32])
...     print(bias.grad.shape)  # torch.Size([1, 1, 2, 128, 128])
torch.Size([1, 128, 2, 128, 32])
torch.Size([1, 128, 2, 128, 32])
torch.Size([1, 128, 2, 128, 32])
torch.Size([1, 128, 2, 128, 32])
torch.Size([1, 1, 2, 128, 128])
On this page