Attention Backends

Our architectures are designed to be agnostic to the attention backend. We implement several attention backends or kernels as described below.

pip install lgatr  # only default attention
pip install lgatr[varlen-attention]  # add varlen attention
pip install lgatr[xformers-attention]  # add xformers attention
pip install lgatr[flex-attention]  # add flex_attention
pip install lgatr[flash-attention]  # add flash attention
pip install lgatr[varlen-attention,xformers-attention,flex-attention,flash-attention]  # add all

The extras install additional requirements, i.e. you don’t have to specify them if you already have these requirements installed. You might have to run python -m pip install --upgrade pip setuptools wheel to update your build environment, extra imports require the most recent versions.

Selecting a backend

Backends are selected per forward call. Every network forwards its extra keyword arguments down to sdp_attention(), so you can name a backend explicitly with backend=...:

outputs, outputs_s = net(multivectors, scalars=scalars, backend="xformers")

If you do not pass backend, the backend is inferred from the other keyword arguments: the block-diagonal mask attn_bias selects xformers, cu_seqlens_q/max_seqlen_q select flash, cu_seq_q/max_q select varlen, and score_mod/block_mask select flex. Without any of them, the native PyTorch backend is used:

from xformers.ops.fmha import BlockDiagonalMask

attn_bias = BlockDiagonalMask.from_seqlens(seqlens)
outputs, outputs_s = net(multivectors, scalars=scalars, attn_bias=attn_bias)

Backends are resolved lazily on first use, so importing lgatr never pulls in xformers or flash-attn. Requesting a backend whose dependency is missing raises a ValueError naming the backend and the reason it could not be loaded.

lgatr.primitives.attention_backends.get_attention_backend(**kwargs)[source]

Resolve the attention backend based on the extra keyword arguments.

Implemented backends:

  • PyTorch native attention: torch.nn.functional.scaled_dot_product_attention

  • PyTorch varlen attention: torch.nn.attention.varlen.varlen_attn

  • xformers attention: xformers.ops.memory_efficient_attention

  • PyTorch flex_attention: torch.nn.attention.flex_attention.flex_attention

  • Flash attention (variable sequence length): flash_attn.flash_attn_varlen_func

The backend is selected explicitly via backend=... if provided, otherwise inferred from backend-specific kwargs (e.g. cu_seqlens_* triggers flash). Falls back to the native backend. Backends are imported lazily on first use.

Return type:

Callable

Why care about Attention Kernels?

As sequence length grows, attention becomes the bottleneck in transformers—both in memory consumption and computation time. This is because attention is the only transformer operation whose cost grows quadratically with the sequence length. To address this, researchers have devoted significant effort to designing more efficient attention backends that reduce this quadratic blowup. The best-known example is FlashAttention (https://arxiv.org/abs/2205.14135), which never writes out the full attention matrix to memory but instead computes attention in smaller chunks. Today, FlashAttention is the standard in most transformer libraries, e.g. torch.nn.functional.scaled_dot_product_attention in PyTorch.

However, these highly optimized CUDA kernels impose constraints on the attention inputs and structure of the attention mask.

  • PyTorch’s default attention path assumes dense tensors for queries, keys, values, and the attention mask. In particle physics applications, each event often contains a different number of particles—so using dense tensors would require padding all events to the maximum length. An alternative is to work with “sparse” representations (e.g., concatenating all particles in a long list and tracking event boundaries), then apply a block-diagonal mask so that only particles within the same event attend to one another. By avoiding operations on padded particles, this approach can dramatically reduce both memory usage and computation time.

  • Beyond sparse masks, many advanced positional-encoding schemes (such as relative positional embeddings, ALiBi, sliding-window attention, PrefixLM, tanh-soft-capping, and so on) also require custom attention kernels that deviate from the dense, full-matrix assumption.

  • Optimized attention backends are typically optimized for the most recent NVIDIA GPUs, and do not or only partially support older hardware.

  • Optimized operations typically only support certain data types. float16 and bfloat16 are most widely supported, and float32 and float8 upcoming.

  • Most backends currently only support certain head dimensions, e.g. only powers of 2.

PyTorch’s native Attention

PyTorch’s native scaled_dot_product_attention is easy to use, but it requires dense tensors for queries, keys and values. It is automatically contained in

pip install lgatr

xformers Attention

Xformers is a library for efficient attention implementations maintained by facebook, including support for block-diagonal attention masks. Unfortunately, xformers does not support MacOS anymore. Under the hood, xformers supports multiple attention backends, including the varlen FlashAttention mentioned below, and selects the best one for your hardware automatically. To the best of our knowledge, the xformers backend is the only attention backend that supports sparse sequence representations and float32. This is the reason why xformers was used in the original L-GATr publications for tasks that require variable-length sequences. To use it, you need to install lgatr with the xformers-attention extra, which requires torch>=2.4:

pip install lgatr[xformers-attention]

PyTorch’s flex_attention

To mitigate the increasing need for custom attention kernels, PyTorch developers have designed flex_attention, a tool that aims to generalize all variations of attention kernels while remaining efficient. The idea is to have two functions score_mod and block_mask as arguments that allow the user to create most attention variants, see this blog. flex_attention is considered stable for torch>=2.7, we therefore do not include it in the default installation of lgatr yet. To install lgatr with flex-attention, run

pip install lgatr[flex-attention]

Official FlashAttention

The original FlashAttention implementation by Tri Dao is still actively maintained and widely used within the community. Its support for variable-length sequences is not yet part of native PyTorch, so we provide an extra for it in lgatr. A problem of this package is that it only supports fp16/bf16 precision on the attention arguments, and we found that this degrades performance in the cases that we tested, see Efficient implementation. Note that xformers might default to using FlashAttention under the hood if it detects that your attention inputs and hardware support it. See its documentation for installation instructions for this package, the process is a bit more involved than for other backends. You can install lgatr with the flash-attention extra, which requires torch>=2.1, as follows:

pip install lgatr[flash-attention]

PyTorch’s native varlen attention

PyTorch 2.10 natively includes a varlen attention kernel that is very similar to the official flash attention implementation. It can be installed with

pip install lgatr[varlen-attention]

More attention backends

L-GATr is designed to be flexible when it comes to attention backends. If you want to use L-GATr with another attention backend, just open an Issue on GitHub, or directly implement it in your own Pull Request!