# attention-gym: examples, masks and kernels for FlexAttention > Attention Gym is a Meta repository of kernels, guides and examples built around PyTorch's FlexAttention, the API for expressing attention variants as composable mask and score modifications rather than as a new kernel. It started as a gallery of examples and has grown into a playground that also ships sparse attention patterns and linear attention implementations. **meta-pytorch/attention-gym** — Helpful tools and examples for working with flex-attention - Repository: https://github.com/meta-pytorch/attention-gym - Stars: 1,262 · Forks: 87 - Language: Python - License: BSD-3-Clause - Published: 2026-10-19 · Updated: 2026-10-19 · Language: en - Canonical page: https://hysenlabs.com/projects/meta-pytorch-attention-gym ## A gallery that grew into a toolkit The overview explains the trajectory, and it explains the shape of the project. Attention Gym began as a library of examples showing the many ways to express attention variants with the FlexAttention API. It is growing into a broader playground for attention, with the addition of sparse attention kernels, linear attention APIs for training and inference, and examples of how to use FlexAttention and friends in real workloads. The starting point matters for understanding the API. FlexAttention is PyTorch's way of expressing attention variants that a standard implementation cannot handle, by letting you write a mask modification and a score modification as Python functions and letting the compiler generate the kernel. A mask says which positions may attend to which, and a score modification changes the value of a score before the softmax. The consequence is that most attention variants become a few lines of Python rather than a new CUDA kernel, which is why a gallery of examples is a genuinely useful artefact rather than a toy. The variants that do not fit the model, the ones with fundamentally different complexity or a different state, are the ones the project has had to add kernels for, and that is what the sparse and linear work is. The feature list is short. FlexAttention masks and score modifications, sparse attention patterns, APIs and kernels for efficient GDN and KDA, utility functions for creating and combining attention masks, and examples of FlexAttention in real-world scenarios. GDN and KDA are the gated delta network style of linear attention models, which is where the newer work is concentrated. The licence is BSD 3-Clause for the code the project authors wrote, and vendored third-party components keep their own licences and notices, which is the usual arrangement for a repository that pulls in other people's kernels. ## Three workflows, and the third is where the real code is The usage section describes three complementary workflows, and the ordering is a hint about which one to start with. The first is composing FlexAttention building blocks. You import the mask modification and score modification functions from the masks and mods modules and pass them directly to PyTorch's FlexAttention APIs. This is the pure-Python path, and the contribution guide asks for exactly this kind of contribution, new masks and score mods, which tells you the library is meant to grow through this surface. The second is using the sparse and linear attention APIs and kernels. The readme names them: a gather attention entry point, the chunk, recurrent and decode paths for GDN and KDA, and short-convolution primitives. The readme points at two working examples, one for compressed sparse attention and one for delta-rule training with the KDA and GDN models. The third is running real workloads and benchmarks. The examples directory covers paged, ring and variable sparse attention, CUDA Graphs, determinism, compilation and profiling. The readme's comment on this is friendly rather than formal, saying most of it should serve as inspiration for fun things you might build from the building blocks. That comment understates the value of that directory. Paged attention and ring attention are the two mechanisms that let a long context sequence be processed in chunks, and they are hard to get right. Determinism and compilation examples are the ones to read if you are going to ship a FlexAttention model, because both are places where a research implementation quietly diverges from a production one. ## One dependency by default, everything else behind extras The installation section is short, and the philosophy in it is worth noticing. ```bash pip install attn-gym ``` The base package intentionally keeps its runtime dependency surface small: it depends only on PyTorch, unpinned. Optional features live behind extras, so you install only what you need: ```bash pip install "attn-gym[linear]" # Linear-attention APIs and kernels ``` ```bash pip install "attn-gym[viz]" # Visualization and example dependencies ``` Only an unpinned PyTorch is a real dependency, and that is a sensible choice for a package whose entire purpose is to exercise the latest features of PyTorch's attention stack. Pinning PyTorch in a library like this would be a support burden with no benefit. The warning that follows the install commands is the most important sentence in the readme. Attention Gym is under active development, the project reserves the right to make backward-incompatible changes between releases, and if you depend on a particular API or kernel behaviour you should hard-pin the version you test, with the example given as pinning the linear extra to a specific version. That is an unusual amount of caution for a version 0.0 release, and it is the right advice. A library of examples and kernels is exactly the kind of thing where the API can move without a major version bump, because the examples are the documentation and the documentation is allowed to change shape. ## The extras tell you which hardware paths exist The extras in the packaging metadata are the clearest description of the project's scope, and they are worth reading as a map. The linear extra pulls in the NVIDIA CUTLASS Python DSL at version 4.6 or newer, but only on Linux, together with the Apache TVM FFI and the packaging library. The cuDNN extra takes the same DSL at 4.7 or newer, again Linux only, and adds the CUDA Python bindings at version 13. So the accelerated linear attention paths are CUDA and Linux, and the cudnn path is a slightly newer DSL than the linear one. The tests extra has pytest and its xdist plugin, typer and jsonargparse for command line interfaces, and a pre-release of FlashAttention 4, again restricted to Linux on x86_64, which tells you the test suite exercises the flash attention path and that it is not expected to run on other platforms. The visualisation extra is the mild one: matplotlib, jsonargparse, tabulate and marimo, and marimo is the reactive notebook environment, which is a good sign for a project whose outputs are visual comparisons. The docs extra is the documentation toolchain: the Material theme for MkDocs, a macros plugin, a redirects plugin and the docstring generator. The dev extra adds the prek hook runner, ruff, pytest and the same CLI libraries, with a comment in the manifest explaining that xdist must be present even when the parallel option is off, because the pytest addopts pass the distribution flag. There is one conflict declared in the tool configuration, between the cudnn extra and the tests extra. That is a real constraint rather than a packaging accident, and it is the kind of detail that saves a day when two extras cannot be installed together. ## The contribution recipe is four steps long The contributing guide asks for contributions, and specifically for new masks or score mods, which is a good sign about where the project wants to grow. The recipe is four steps. Create a new file in the masks directory for a mask modification, or in the mods directory for a score modification. Implement your function and add a simple main function that showcases it. Update the init file of the module to include your new function. Optionally, add an end to end example using it in the examples directory. Three things about that recipe are worth noticing. The first is the showcase main function. Every contribution has to be runnable, and the simplest way to make that a habit is to require a demonstration at the bottom of the file. It also means the library stays runnable as a gallery, which is what it was. The second is that the index file is updated by hand. For a library that could generate it, that is a deliberate choice, probably so that the export order and the grouping are the author's rather than alphabetical. The third is that an example is optional. The guide says optionally, add an end to end example, which is the difference between a small contribution and a larger one, and it matches the third workflow in the usage section. The rest of the guide is in the contributing document, and the project has the other standard pieces, a design document, a code of conduct, an agent instructions file and a second instructions file aimed at coding assistants, plus a directory of agent configuration. ## The repository layout, and the hooks that keep it honest The tree shows a project with an unusually broad surface. There is the package itself, a benchmarks directory, a docs directory, an examples directory split into flex attention, linear and sparse groups, a test directory, a tools directory, a main entry point at the root, and a modal test file, which suggests some of the heavier tests run on rented GPU instances rather than on a laptop. There is a Makefile-free layout with the packaging in a pyproject file, built with hatchling and hatch-vcs, where the version comes from the version control tag rather than being written in the file. The developer workflow is short. Install the development requirements in editable mode with the dev extra, then install and run the repository hooks: ```bash pip install -e ".[dev]" ``` ```bash prek install prek run --all-files ``` The hook runner is prek, which is the Rust reimplementation of the pre-commit framework, and running it across all files is a pre-commit check without the Python dependency. The version history is a nice illustration of what a young kernel library looks like. The most recent release, 0.0.17, is titled deterministic l2norm. The one before, 0.0.16, is a Triton 3.9 fix for the GDN backward pass. The one before that, 0.0.15, is a CuTeDSL 4.6.0 update for the linear path together with an update to the TVM FFI. Read as a sequence, those release titles are a maintenance log for a project that is tracking the fast moving parts of the GPU stack: a Triton version, a CUTLASS DSL version, an FFI version. That is the cost of building on top of them, and it is also the reason the project tells you to pin. ## Conclusion attention-gym suits a PyTorch engineer who wants to try an attention variant without writing a Triton kernel, because the mask and score modification API lets you express most sparsity and bias patterns in a few lines of Python. It is a poor fit if you need a stable API, since the readme warns that backward-incompatible changes are reserved between releases, and a poor fit if your hardware is not the one the extras target, since the linear and cuDNN paths are Linux and CUDA only. Before you depend on it, read the warning about pinning versions, check which extras your platform can actually install, and treat the examples directory as the documentation, since it is where the working patterns live. ## FAQ ### How do I install attention-gym? With pip install attn-gym. The base package depends only on PyTorch, unpinned, and optional features live behind extras, so you can install attn-gym[linear] for the linear attention APIs and kernels or attn-gym[viz] for the visualisation dependencies. PyTorch 2.5 or higher and Python 3.10 or higher are required. ### What are mask_mod and score_mod in attention-gym? They are the two kinds of building block you pass to PyTorch's FlexAttention APIs. Mask modifications say which positions may attend to which, and score modifications change a score before the softmax. The library ships collections of both, and the contribution guide asks for new ones, with a four step recipe for adding one. ### Is attention-gym stable enough to depend on? The readme warns that it is under active development and that the project reserves the right to make backward-incompatible changes between releases. If you rely on a particular API or kernel behaviour, hard-pin the version you test, for example by installing the extra with an exact version. ### Which platforms can install the accelerated extras? The linear and cuDNN extras require the NVIDIA CUTLASS Python DSL and are marked for Linux only, and the tests extra includes a pre-release of FlashAttention 4 restricted to Linux on x86_64. The cudnn extra and the tests extra are also declared as conflicting, so they cannot be installed together. ## Sources - [Issues](https://github.com/meta-pytorch/attention-gym/issues) - [License: BSD-3-Clause](https://github.com/meta-pytorch/attention-gym/blob/main/LICENSE) - [meta-pytorch/attention-gym on GitHub](https://github.com/meta-pytorch/attention-gym) - [README](https://github.com/meta-pytorch/attention-gym/blob/main/README.md) - [Releases](https://github.com/meta-pytorch/attention-gym/releases) --- Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/meta-pytorch-attention-gym