XFormers

Warnings

• breaking Strict compatibility requirements with PyTorch and CUDA versions. Installing 'xformers' via pip without specifying a PyTorch index URL can lead to incompatibility issues or an unwanted PyTorch upgrade.
  Fix: Always install xformers with a matching PyTorch wheel via `--index-url https://download.pytorch.org/whl/cuXXx` (e.g., `cu126`) or build from source if a specific PyTorch/CUDA combination is needed.
• gotcha Many xFormers optimizations, particularly `memory_efficient_attention`, can produce non-deterministic results, meaning repeated runs with the same inputs might yield slightly different outputs.
  Fix: If reproducibility is critical, avoid using xFormers' non-deterministic kernels or use alternative attention implementations. Check `xformers.info` output for deterministic kernel availability.
• breaking Dropped support for V100 and older NVIDIA GPUs, following PyTorch's deprecation schedule. Flash-Attention 2 support for building as part of xFormers is also deprecated.
  Fix: Upgrade GPU hardware or use older xFormers/PyTorch versions if V100 compatibility is essential. Consider using Flash-Attention 3 on Ampere GPUs or Flash-Attention 2 through PyTorch on Linux.
• deprecated Many classes and modules within `xformers.factory`, `xformers.triton`, and `xformers.components` have been or will be deprecated.
  Fix: Consult the official xFormers documentation and CHANGELOG for updated API usage and recommended alternatives for constructing Transformer components.
• breaking The `memory_efficient_attention` function now expects the `attn_bias` argument to explicitly have a head dimension. It no longer automatically broadcasts batch/head dimensions for `attn_bias`.
  Fix: Manually ensure your `attn_bias` tensor has the correct dimensions, including the head dimension (e.g., by using `.expand()` or `.repeat()` if broadcasting is intended).
• gotcha Building xFormers from source on Windows can be complex due to dependencies on Visual Studio Build Tools, specific CUDA Toolkit versions, and potential long path issues. Pre-built wheels are highly recommended.
  Fix: Prioritize `pip install` with appropriate PyTorch/CUDA index URLs. If building from source, ensure Visual Studio Build Tools (C++ desktop development), correct CUDA Toolkit, and `git config --global core.longpaths true` are configured.

Install

• pip install xformers Latest stable release (requires latest PyTorch)
• pip install -U xformers --index-url https://download.pytorch.org/whl/cu126 For CUDA 12.6 (adjust URL for other CUDA versions like cu118, cu128, cu130)
• pip install --pre -U xformers Latest development binaries
• pip install ninja pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers From source (for specific PyTorch versions or nightlies)

Imports

• memory_efficient_attention

  from xformers.ops import memory_efficient_attention

  The primary optimized attention function.
• LowerTriangularMask

  from xformers.ops.fmha.attn_bias import LowerTriangularMask

  Common attention bias for causal masking.
• AttentionOpBase

  from xformers.ops import AttentionOpBase

  Base class for attention operators, used for dispatching or enforcing specific backends.

Quickstart

This quickstart demonstrates how to use `xformers.ops.memory_efficient_attention` with dummy PyTorch tensors for both standard and causal attention patterns. It highlights the typical tensor shape and the common practice of using half-precision floating-point numbers (float16) for performance on GPUs. The `xformers.info` utility is also mentioned for diagnostics. Ensure PyTorch and CUDA are properly installed and configured.

import torch
from xformers.ops import memory_efficient_attention, LowerTriangularMask

# Ensure tensors are on CUDA if available
device = "cuda" if torch.cuda.is_available() else "cpu"

# Assume batch_size=2, seq_len=128, num_heads=8, head_dim=64
batch_size = 2
seq_len = 128
num_heads = 8
head_dim = 64

# Create dummy query, key, value tensors
# xFormers memory_efficient_attention typically expects (batch_size, seq_len, num_heads, head_dim)
query = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device)
key = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device)
value = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device)

# It's common to use float16 (half precision) for performance with xFormers
query = query.half()
key = key.half()
value = value.half()

# Example 1: Standard memory-efficient attention
# xFormers automatically dispatches to the best available operator
output_attn = memory_efficient_attention(query, key, value)
print(f"Output attention shape (standard): {output_attn.shape}")

# Example 2: Causal attention with a lower triangular mask
# Note: The attn_bias argument structure has changed in newer versions (e.g., v0.0.21+)
# For LowerTriangularMask, it often handles internal expansion if num_heads is implicitly available.
attn_bias = LowerTriangularMask()
output_causal_attn = memory_efficient_attention(query, key, value, attn_bias=attn_bias)
print(f"Output attention shape (causal): {output_causal_attn.shape}")

# To verify installation and available kernels:
# import subprocess
# subprocess.run(["python", "-m", "xformers.info"])