# MagiAttention: distributed attention for very long context

> MagiAttention is SandAI's Apache-2.0 library for attention whose cost scales linearly with context length, shipped as CUDA extensions alongside training examples for Megatron, native PyTorch and Transformers. It targets distributed training on multi-GPU clusters with NVSHMEM, and it is not a drop-in replacement for standard attention in a small single-GPU setup.

**SandAI-org/MagiAttention** — A Distributed Attention Towards Linear Scalability for Ultra-Long Context, Heterogeneous Data Training

- Repository: https://github.com/SandAI-org/MagiAttention
- Stars: 971 · Forks: 74
- Language: Python
- License: Apache-2.0
- Published: 2026-10-20 · Updated: 2026-10-20 · Language: en
- Canonical page: https://hysenlabs.com/projects/sandai-org-magiattention

## The goal is linear cost as context grows

MagiAttention is a distributed attention library from SandAI, and its stated aim is a form of attention whose cost scales linearly for ultra-long context, with support for heterogeneous mask training. The project name, the package name magi_attention and the description in the packaging metadata all use the same phrasing, and a paper is referenced at arxiv 2505.13211.

The problem it targets is the one that shows up in long context pretraining. Attention cost grows with the square of sequence length, so a team that wants to train on a million tokens of context is not short of compute, they are short of a way to make the attention itself affordable. A distributed formulation that spreads the attention computation across devices is the standard answer, and MagiAttention's angle is to combine that distribution with mask handling that is not uniform across the sequence.

That mask dimension is what the heterogeneous data training wording refers to. When a training batch mixes documents of different lengths, or packs sequences with padding, the attention pattern is no longer a single uniform structure, and a distributed implementation has to carry that variation without falling back to a dense computation.

The audience is therefore training infrastructure engineers and researchers working on long context models, on hardware clusters rather than laptops. The distribution of releases and the presence of NVSHMEM in the dependency list both point that way.

## A Python package wrapped around compiled CUDA sources

The repository top level shows a project with two halves. On one side there is magi_attention/, the importable Python package, and on the other there is the machinery that compiles kernels for it.

Inside the package, the packaging configuration includes the common headers and utilities for runtime and JIT compilation, the extension sources, and a directory of flexible flash attention sources with .h, .hpp, .cuh, .cu and .cpp files. The package also ships .pyi stub files, which means type information is provided for a codebase whose internals are largely C++ and CUDA.

At the top level of the repository, extensions/ holds the compiled pieces, exps/ holds experiments, tools/ holds utilities, scripts/ holds helper scripts, tests/ and conftest.py hold the test suite, docs/ holds documentation, and assets/ holds images used by the README. There is also a .gitmodules file, so parts of the tree arrive as git submodules, a .pre-commit-config.yaml, a mypy.ini, a .flake8 and a .clang-format.

Two build directories at the top level are worth noting because they tell you what integration targets are supported: examples/megatron/, examples/torch_native/ and examples/transformers/. That trio is the practical answer to the question of how you would actually use this, and it is covered in more detail below.

## The build compiles CUDA extensions with ninja and pinned NVIDIA packages

There is no wheel for this project as far as the repository shows. setup.py is a real build script, not a shim, and it imports the CUDA extension machinery from PyTorch directly:

```python
from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension
```

The declared build requirements confirm the shape of the build:

```toml
requires = [
    "setuptools>=61.0",
    "wheel",
    "torch",
    "ninja",
    "packaging"
]
```

Ninja is in that list because it is the build backend used to compile the CUDA sources, and torch is there because the extension is built as a PyTorch CUDA extension. A .clang-format file at the top level indicates the C++ sources are formatted with clang-format, and the .coveragerc plus .flake8 and mypy.ini show that both the C++ and Python sides are linted.

The runtime requirements are pinned tightly, which is a deliberate choice for a kernel library:

```bash
versioningit
packaging>=24.2
setuptools>=78.1.1
nvidia-nvshmem-cu12==3.4.5
nvidia-cutlass-dsl==4.4.2
apache-tvm-ffi==0.1.8.post0
quack-kernels==0.4.1
```

NVSHMEM is the collective communication library for the distributed path, CUTLASS is the template library for the kernels, and quack-kernels and apache-tvm-ffi come in as pinned dependencies. Note the cu12 suffix on the NVSHMEM package, which ties this build to a CUDA 12 generation environment.

Two environment variables in setup.py affect the build. One locates the project home:

```python
USER_HOME = os.getenv("MAGI_ATTENTION_HOME")
```

And one forces a particular C++ ABI setting, described in a comment as needed for CI images that use the C++11 ABI:

```python
FORCE_CXX11_ABI = os.getenv("MAGI_ATTENTION_FORCE_CXX11_ABI", "0") == "1"
```

The script also handles a CUDA 13 specific detail, adding the CCCL include path explicitly for CUDA13 and later, and it warns that ninja builds need absolute include directories.

## Heterogeneous mask training and the distributed roll API

The README news section records the most recent feature work, and the entry for version 1.1.1 introduces a distributed roll API for MTP, which is the multi-token prediction setup used in recent large model training. The entry is cut off in the visible README text, so the API surface is not documented there, but the fact of the introduction is the relevant part for anyone planning to adopt a specific version.

The roll API deserves attention because it names a shape of work rather than an optimisation. A roll operation in this context is a forward pass that produces predicted tokens and keeps the state needed to continue, and doing that in a distributed setting means the coordination of partial results across ranks. Introducing a distributed version of it is a signal that the project is moving beyond attention primitives towards the training loop itself.

The heterogeneous mask training thread runs through the whole project rather than one feature. The project tagline pairs linear scalability with heterogeneous mask training, and the packaged flexible flash attention sources are consistent with an implementation that can vary its attention pattern per sample or per segment.

For an adopter, this is the part to watch in the release notes rather than to expect from the README. The project publishes docs at its documentation site, a blog describing the approach, and a release feed, and those are where the API details live.

## Three examples show the intended integration paths

The examples/ directory is the most honest description of how MagiAttention is meant to be used, and it contains three subdirectories: megatron/, torch_native/ and transformers/.

The Megatron example targets the parallel training framework that already has its own sequence parallelism mechanisms, so this example is about plugging the attention implementation into an existing distributed training stack. The transformers example targets the Hugging Face model library, which is where a team that wants long context in an existing transformer codebase would start. The torch_native example is the plain PyTorch path, which is the way to understand the API without a training framework in the way.

The presence of all three tells you the project is not a single integration. It is a kernel and dispatch layer that has to be adapted to whichever framework owns your training loop, and the effort of adoption is therefore proportional to how much of your stack you are willing to change.

The experiments directory at the top level, exps/, is separate from examples, which suggests the distinction the project draws is between code that demonstrates the API and code that reproduces training runs. The tests/ directory with conftest.py at the root is where the behaviour is checked.

Note also that the tool is distributed under the Apache-2.0 licence, with the LICENSE file at the root of the repository and the licence declared in the packaging metadata as well, and the copyright header in setup.py names SandAI as the holder.

## Versions come from git tags through versioningit

The package does not hardcode a version. The packaging metadata marks the version as dynamic and configures versioningit to read it from the vcs, with a default version of 1.1.6 and a write target that writes the resolved version into magi_attention/_version.py using the template `__version__ = '{version}'`.

In practice that means two things. Installing from a checkout gives you the version of the tag or commit you are on, and the build writes a version file inside the package rather than relying on a hardcoded constant. The default-version setting acts as a fallback when no tag information is available, which matters for source archives built outside a git checkout.

The release history is close and regular. Version 1.1.6 was tagged on 2026-09-18, with 1.1.5 and 1.1.4 both tagged on 2026-09-15. That pattern suggests small, frequent releases rather than a long schedule, which suits a project whose kernels are being adapted to changing frameworks.

The last push to the main branch was on 2026-10-18, so work is continuing after the 1.1.6 tag. If you adopt this, pin a tag rather than tracking main, and expect the gap between a tag and the branch head to contain changes you have not evaluated.

The project also names its authors, Zewei Tao and Yunpeng Huang, and points bug reports at the repository issue tracker, with a bug tracker URL in the packaging metadata and a contributing guide at the root.

## Where MagiAttention is the wrong tool

The clearest case for not using it is a context length that fits on one device. If your sequence fits comfortably in the memory of a single GPU, the distributed attention machinery has nothing to distribute, and you are left with a build that compiles CUDA extensions, requires an NVSHMEM enabled environment and adds a dependency chain of pinned NVIDIA packages, in exchange for nothing. The plain PyTorch attention path inside your existing framework is the better tool in that situation.

The second constraint is the environment. The pinned nvidia-nvshmem-cu12 build, the ninja driven CUDA compilation, the CUDA 13 CCCL include handling and the C++11 ABI escape hatch all point to a build that assumes a specific cluster setup. Reproducing it on a workstation or on a managed notebook service is a different exercise from using it where it was designed to run, and the README text visible here does not promise container images or prebuilt wheels to sidestep that.

The third boundary is the risk profile of a kernel library. The Python API can change, but a compiled kernel has to match the PyTorch version, the CUDA version and the framework you plug it into. A framework upgrade in your training stack is where this dependency will break first, and the project does not publish a compatibility matrix for that.

The alternative approach for a team that needs long context but not distributed attention is to keep the standard attention implementation and reduce the cost elsewhere, for instance with sequence parallel training in your framework of choice, and reserve a specialised kernel for the layers where it actually matters. The difference in approach is that you own the coordination and the maintenance, and you avoid a compiled dependency, at the cost of doing the distributed work yourself.

## Conclusion

Adopt MagiAttention when you are training very long context models on a multi-GPU cluster with NVSHMEM available and the standard sequence parallel attention path is the bottleneck, and skip it when your context fits comfortably in one device, where the distributed machinery buys nothing. Verify three things before you commit: that the pinned nvidia-nvshmem-cu12 build on your cluster matches the CUDA version you have, that the submodules under .gitmodules are initialised, since the CUDA sources live partly outside the Python package, and that the Megatron or torch_native example under examples/ matches your training stack before you port the attention op into it. The last push to main was on 2026-10-18, version 1.1.6 was tagged on 2026-09-18, and the project version is derived from git tags by versioningit, so pin a tag rather than tracking main.

## FAQ

### What is MagiAttention used for?

MagiAttention is a distributed attention implementation aimed at linear scalability for ultra-long context, with support for heterogeneous mask training. The project describes itself as a distributed attention for ultra-long context and heterogeneous data training.

### What are the installation requirements for MagiAttention?

Python 3.10 or later is required, and the build compiles CUDA extensions using setuptools, torch, wheel, ninja and packaging as build requirements. The runtime requirements pin nvidia-nvshmem-cu12, nvidia-cutlass-dsl, apache-tvm-ffi and quack-kernels to exact versions.

### How does MagiAttention integrate with training frameworks?

The examples directory contains separate examples for Megatron, native PyTorch and Transformers, so the attention implementation is adapted to whichever framework owns your training loop rather than shipped as one fixed integration.

### Which version of MagiAttention should I pin?

Pin a release tag rather than the main branch. The version is derived from git tags by versioningit with a default of 1.1.6, and 1.1.6 was tagged on 2026-09-18 while the last push to main was on 2026-10-18.

### What licence is MagiAttention released under?

MagiAttention is released under the Apache-2.0 licence. The LICENSE file sits at the root of the repository and the packaging metadata declares the same licence.

## Sources

- [Issues](https://github.com/SandAI-org/MagiAttention/issues)
- [License: Apache-2.0](https://github.com/SandAI-org/MagiAttention/blob/main/LICENSE)
- [README](https://github.com/SandAI-org/MagiAttention/blob/main/README.md)
- [Releases](https://github.com/SandAI-org/MagiAttention/releases)
- [SandAI-org/MagiAttention on GitHub](https://github.com/SandAI-org/MagiAttention)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/sandai-org-magiattention