Installation#
Warning
MagiAttention currently supports only Hopper and Blackwell. We are actively working to support more GPU architectures in upcoming releases.
Setup Environment#
Activate an NGC-PyTorch Container#
Tip
We recommend you to use the standard NGC-PyTorch Docker Releases for consistency of basic dependencies such as Python, CUDA, PyTorch, etc.
Warning
Due to performance issue caused by CUDA-12, we recommend you to use CUDA-13+ based NGC-PyTorch containers for optimal performance.
And we add an assertion in the setup.py script to check the CUDA version and abort the installation if the CUDA version is lower than 13.0.
If you insist on using CUDA-12 based containers, you can set the environment variable MAGI_ATTENTION_ALLOW_BUILD_WITH_CUDA12=1, but please be aware that it may lead to significant performance degradation compared to CUDA-13+.
docker run command:
# choose one compatible version, e.g. 25.10, 25.12, etc. MAJOR_VERSION=25 MINOR_VERSION=10 # specify your own names and paths CONTAINER_NAME=... HOST_MNT_ROOT=... CONTAINER_MNT_ROOT=... docker run --name ${CONTAINER_NAME} -v ${HOST_MNT_ROOT}:${CONTAINER_MNT_ROOT} -it -d --privileged --gpus all --network host --ipc host --ulimit memlock=-1 --ulimit stack=67108864 nvcr.io/nvidia/pytorch:${MAJOR_VERSION}.${MINOR_VERSION}-py3 /bin/bash
docker exec command:
docker exec -it ${CONTAINER_NAME} /bin/bash
Pull Source Code#
git commands:
git clone https://github.com/SandAI-org/MagiAttention.git cd MagiAttention git submodule update --init --recursive
Enable IBGDA (optional)#
Note
If you would like to try using our native group-collective kernels when cp_size > 8 as the communication backend, i.e. a process group involving both intranode (connected through NVLink) and internode (visible through RDMA) peers, you’re required to enable IBGDA on your bare-metal host machine.
Warning
This step needs to be performed on the BARE-METAL HOST OPERATING SYSTEM, NOT inside a Docker or other containerized environment, as containers do not manage the host kernel.
bash script:
bash scripts/enable_ibgda_on_host.sh
Setup Dependencies#
Install Required Packages#
pip install command:
pip install -r requirements.txt
Note
requirements.txt accepts nvidia-cutlass-dsl 4.4.2 to 4.8.1 except 4.6.0 and 4.6.1, and quack-kernels 0.4.1 to 0.6.5. The releases 4.6.0 and 4.6.1 add their own thread election around non-tensor bulk copies. The MagiAttention kernels already wrap those copies in elect_one, so the two elections nest. That combination was not run and the two releases are excluded.
Each quack-kernels release calls nvidia-cutlass-dsl APIs that other DSL releases removed. The supported pairs are:
|
|
|---|---|
4.4.2 |
0.4.1 |
4.6.2 |
0.6.4 |
4.7.0 to 4.8.1 |
0.6.5 |
The runtime check is stricter than pip. pip accepts any combination inside the two requirements.txt ranges, but import magi_attention checks the installed pair against this table when it loads the CuTeDSL kernels. If the pair is not in the table, or one of the two packages is not installed, the import raises CuteDslRequirementError for every kernel backend. This command installs the newest supported pair:
pip install --upgrade "nvidia-cutlass-dsl[cu13]<=4.8.1" "quack-kernels<=0.6.5"
A fresh pip install -r requirements.txt resolves to the newest pair, 4.8.1 and 0.6.5. In an existing environment, upgrade both packages together with the command above. Do not use pip install --upgrade -r requirements.txt for this, because it also upgrades unrelated packages such as setuptools and packaging.
CI tests the two ends of the range: 4.4.2 with 0.4.1, and 4.8.1 with 0.6.5. The two middle pairs were tested only outside CI, on Blackwell. 4.6.2 with 0.6.4 passed tests/test_kernel/cutedsl/ and the cutedsl cases of tests/test_pipeline.py with head dim 64 and 128 at world size 1 and 2. 4.7.0 with 0.6.5 passed tests/test_kernel/cutedsl/. Other releases inside the range were not run.
Warning
With nvidia-cutlass-dsl below 4.8.0, a compiled CuTeDSL function can hang when three or more threads of one process call it for the first time at the same time (CUTLASS #3706). This was reproduced with 4.4.2, 4.6.2 and 4.7.0, and not with 4.8.0. Two threads, or a first call from one thread before the other threads call the function, do not hang. MagiAttention does not call a CuTeDSL function from several threads at the same time.
Install flash_attn_cute (optional)#
Note
If you want the FFA_FA4 backend on Ampere or Blackwell, install flash_attn_cute. The default ffa backend supports Hopper only.
Warning
The FFA_FA4 backend (flash_attn_cute, built from the FlashAttention submodule) requires nvidia-cutlass-dsl>=4.4.2,<4.6.0. It uses cute.core.ThrMma and cute.make_fragment, which 4.6.0 removed. The upper bound can be lifted when the FlashAttention submodule moves to a flash_attn_cute release that uses neither API. The package metadata of flash_attn_cute declares only the lower bound, so pip does not enforce the upper bound. On 4.6.0 and newer, flash_attn_cute fails to import and MagiAttention treats the fa4 backend as not installed. The other backends keep working if the installed nvidia-cutlass-dsl / quack-kernels pair is in the table of supported pairs above. If you select the fa4 backend, MagiAttention raises a RuntimeError that names the installed nvidia-cutlass-dsl version, the required range, the original import error, and the commands below.
FFA_FA4 was tested with nvidia-cutlass-dsl 4.4.2 only. To use it, install the pair 4.4.2 and 0.4.1, which is also a supported pair for the other backends:
pip install "nvidia-cutlass-dsl[cu13]==4.4.2" "quack-kernels==0.4.1"
Then install or reinstall flash_attn_cute from the MagiAttention source root with the bash script below.
bash script:
ARCHS="sm80,sm100" # if you only want to use in Blackwell, you can set it to "sm100" to speed up the installation; and if you only want to use in Ampere, you can set it to "sm80" accordingly bash scripts/install_flash_attn_cute.sh $ARCHS
Install MagiAttention#
Install MagiAttention From Source#
Warning
This progress may take around 10~20 minutes and occupies 90% of CPU resources for the first time.
Note
We have several environment variables to fine-grained control the installation progress, especially for CUDA extension modules building.
pip install command for Hopper:
pip install --no-build-isolation .
pip install command for Blackwell:
export MAGI_ATTENTION_PREBUILD_FFA=0 # for now, native ffa does not support Blackwell pip install --no-build-isolation . export MAGI_ATTENTION_KERNEL_BACKEND=fa4 # requires flash_attn_cute, see "Install flash_attn_cute" above
pip install command for Ampere:
export MAGI_ATTENTION_PREBUILD_FFA=0 # for now, native ffa does not support Blackwell export MAGI_ATTENTION_SKIP_MAGI_ATTN_COMM_BUILD=1 # for now, magi_attn_comm does not support Ampere pip install --no-build-isolation . export MAGI_ATTENTION_KERNEL_BACKEND=fa4 # requires flash_attn_cute, see "Install flash_attn_cute" above
PreCompile FFA_FA4 kernels (optional)#
Note
If you would like to try MagiAttention on Blackwell and you’ve already installed both magi_attention and flash_attn_cute to enable FFA_FA backend, we further recommend you to pre-compile the common cases for FFA_FA4 kernels before production usage to avoid runtime JIT re-compilation overhead, since it is built upon Cute PythonDSL.
And the cache directory for pre-compiled kernels is /path/to/magi_attention/lib/ffa_fa4_cache/ by default, which can be overridden by setting the environment variable MAGI_ATTENTION_FFA_FA4_CACHE_DIR to specify a custom cache directory if needed.
python script:
# You can change the cases to pre-compile in the script according to your needs, # and the whole pre-compilation progress will be richly logged # in the terminal by tqdm, for you to track the progress and results. # Parallel worker processes compile the kernels. Compilation uses mostly the CPU. # Use `--num-workers` to set the number of workers (default: min(#CPUs, 64)). python tools/precompile_ffa_fa4.py