DOCSEARCHbyAlgolia
Submit Docs

xFormers Documentation

repository·main·Indexed 27 days ago

https://github.com/facebookresearch/xformers

A toolbox of customizable, high-performance building blocks designed to accelerate Transformer research. xFormers provides memory-efficient and fast CUDA kernels for components such as attention, layer norm, and fused operations. It includes specialized implementations of Fused Multi-Head Attention (FMHA) via backends like Cutlass, Flash, and CK, as well as tools for generating axial, local, and Swin Transformer attention patterns.

Tokens
3K
Snippets
12
Records
21
Agent score
95%

What's inside xFormers

  1. Overview of xFormers

    main
    xFormers is a toolbox designed to accelerate research on Transformers. It provides a collection of interoperable and optimized building blocks that are field-agnostic, allowing them to be used across various domains. The library is designed to be composable (allowing for the creation of a 'block zoo' for architecture search and ablations), extensible (easy to extend locally for specific improvements), and highly optimized through heavy benchmarking and automated testing of all variants.

  2. Key xFormers optimized components

    main

    xFormers provides several optimized building blocks that go beyond standard PyTorch primitives, including:

    • Memory-efficient exact attention
    • Sparse attention
    • Block-sparse attention
    • Fused softmax
    • Fused linear layer
    • Fused layer norm
    • Fused dropout(activation(x+bias))
    • Fused SwiGLU

  6. Install xFormers from source

    main

    If you need to use a specific version of PyTorch (including nightly releases), you can build from source.

    Prerequisites & Tips:

    • PyTorch must already be installed before running the build.
    • Installing ninja can make the build much faster.
    • Use the TORCH_CUDA_ARCH_LIST environment variable if you are building for different GPU architectures than the one currently running.
    • The build process can take dozens of minutes.
    # (Optional) Makes the build much faster
pip install ninja

# Set TORCH_CUDA_ARCH_LIST if running and building on different GPU types
# NOTE: pytorch must already be installed!
pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers

  7. Install xFormers via pip (Recommended)

    main

    For Linux and Windows, the recommended way to install the latest stable version of xFormers is via pip. This requires PyTorch 2.10.0. Choose the command corresponding to your CUDA version.

    # [linux & win] cuda 12.6 version
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu126

# [linux & win] cuda 12.8 version
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu128

# [linux & win] cuda 13.0 version
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu130

# [linux only] (EXPERIMENTAL) rocm 7.1 version
pip3 install -U xformers --index-url https://download.pytorch.org/whl/rocm7.1

  8. Implement Swin Transformer using xFormers sparse kernels

    main

    To implement a Swin Transformer using xFormers' efficient sparse kernels, follow these steps:

    1. Generate the attention mask using AP.swin_attention_pattern.
    2. Convert the mask into a sparse representation using xformers.components.attention._sputnik_sparse.SparseCS.
    3. Use xformers.components.attention.core.scaled_dot_product_attention within your attention module to perform the computation.

    This approach allows you to replicate Swin Transformer behavior without writing custom window-shifting logic, as the complexity is handled by the sparsity pattern.

    import torch
import xformers.components.attention.attention_patterns as AP
from xformers.components.attention.core import scaled_dot_product_attention
from xformers.components.attention._sputnik_sparse import SparseCS

# 1. Create the pattern
attn_mask = AP.swin_attention_pattern(input_resolution[0], input_resolution[1], window_size, shift_size=shift_size)

# 2. Convert to SparseCS for efficient kernels
attn_mask = SparseCS(attn_mask, torch.device("cuda"))

# 3. Use in scaled_dot_product_attention
# x = scaled_dot_product_attention(q, k, v, attn_mask, dropout=attn_drop)

  9. Troubleshoot xFormers installation

    main

    If you encounter issues during installation or build, check the following:

    • CUDA/NVCC Mismatch: Ensure NVCC and the current CUDA runtime match. You may need to adjust your environment using module load commands.
    • GCC Version: Ensure your GCC version is compatible with your NVCC capabilities.
    • GPU Architectures: Set the TORCH_CUDA_ARCH_LIST environment variable to the architectures you want to support. A comprehensive setup is: export TORCH_CUDA_ARCH_LIST="6.0;6.1;6.2;7.0;7.2;7.5;8.0;8.6".
    • Out of Memory (OOM) during build: Reduce parallelism by setting the MAX_JOBS environment variable (e.g., MAX_JOBS=2).
    • Windows Long Paths: If you see Filename longer than 260 characters, enable long paths in Windows and run git config --global core.longpaths true.

  12. Run other FMHA inference benchmarks on ROCM

    main

    Use the following scripts to run additional FMHA inference benchmarks on ROCM hardware:

    • benchmark_attn_decoding.py: For attention decoding benchmarks.
    • benchmark_mem_eff_attention_mqa.py: For memory-efficient multi-query attention (MQA) benchmarks.
    python xformers/benchmarks/benchmark_attn_decoding.py
python xformers/benchmarks/benchmark_mem_eff_attention_mqa.py