Attention Gym

Attention Gym is a collection of helpful tools and examples for working with flex-attention

📚 Docs | 🎯 Features | 🚀 Getting Started | 💻 Usage | 🛠️ Dev | 🤝 Contributing | ⚖️ License

📖 Overview

This repository aims to provide a playground for experimenting with various attention mechanisms using the FlexAttention API. It includes implementations of different attention variants, performance comparisons, and utility functions to help researchers and developers explore and optimize attention mechanisms in their models.

🎯 Features

• Implementations of various attention mechanisms using FlexAttention
• Reference implementations for linear-attention recurrences
• Utility functions for creating and combining attention masks
• Examples of how to use FlexAttention in real-world scenarios

🚀 Getting Started

Prerequisites

• PyTorch (version 2.5 or higher)

Installation

git clone https://github.com/meta-pytorch/attention-gym.git
cd attention-gym
pip install .

💻 Usage

There are two main ways to use Attention Gym:

1. Run Example Scripts: Many files in the project can be executed directly to demonstrate their functionality:

   python attn_gym/masks/document_mask.py
   

   These scripts often generate visualizations to help you understand the attention mechanisms.

2. Import in Your Projects: You can use Attention Gym components in your own work by importing them:

   from torch.nn.attention.flex_attention import flex_attention, create_block_mask
   from attn_gym.masks import generate_sliding_window
   
   # Use the imported function in your code
   sliding_window_mask_mod = generate_sliding_window(window_size=1024)
   block_mask = create_block_mask(sliding_window_mask_mod, 1, 1, S, S, device=device)
   out = flex_attention(query, key, value, block_mask=block_mask)
   

For comprehensive examples of using FlexAttention in real-world scenarios, explore the examples/ directory. These end-to-end implementations showcase how to integrate various attention mechanisms into your models.

Note

Attention Gym is under active development, and we do not currently offer any backward compatibility guarantees. APIs and functionalities may change between versions. We recommend pinning to a specific version in your projects and carefully reviewing changes when upgrading.

📁 Structure

Attention Gym is organized for easy exploration of attention mechanisms:

🔍 Key Locations

• attn_gym.masks: Examples creating BlockMasks
• attn_gym.mods: Examples creating score_mods
• attn_gym.linear: Reference linear-attention implementations
• examples/: Detailed implementations using FlexAttention
• examples/paged_attention: Paged KV-cache examples using FlexAttention

🛠️ Dev

Install dev requirements

pip install -e ".[dev]"

Install and run the repository hooks:

prek install
prek run --all-files

🤝 Contributing

We welcome contributions to Attention Gym, especially new Masks or score mods! Here's how you can contribute:

Contributing Mods

1. Create a new file in the attn_gym/masks/ for mask_mods or attn_gym/mods/ for score_mods.
2. Implement your function, and add a simple main function that showcases your new function.
3. Update the attn_gym/*/__init__.py file to include your new function.
4. Optionally, add an end-to-end example using your new function in the examples/ directory.

See CONTRIBUTING.md for more details.

⚖️ License

attention-gym is released under the BSD 3-Clause License.