facebookresearch/xformers v0.0.36

v0.0.36: fmha moves to MSLK

on GitHub

Improved

• The open-source build no longer compiles anything: xFormers is now pure Python. Everything that needs a kernel lives in separate packages, chiefly mslk. Building from source no longer needs a CUDA toolchain, and no longer needs the CUTLASS submodule.
• xFormers now ships as a single py3-none-any wheel, instead of one wheel per platform and CUDA/ROCm version, so there is no longer any need to match the wheel to your PyTorch build. It is published to all the same places as before — PyPI and each of the PyTorch indices — so none of the installation commands documented in the README need to change. The same wheel also installs on macOS, but the attention ops aren't available there, as mslk has no macOS build.
• The single pure-Python wheel expands platform support: xFormers is now also available for Linux on ARM (aarch64), and on the cu129, cu132, cu134, rocm7.0, rocm7.2 and rocm7.14 PyTorch indices, i.e., on every CUDA and ROCm index of a PyTorch version that xFormers supports.
• Selective activation checkpointing (xformers.checkpoint) now delegates to PyTorch's public torch.utils.checkpoint.create_selective_checkpoint_contexts instead of reaching into PyTorch private internals. This fixes an IndexError with recent PyTorch versions. The public API (checkpoint, get_optimal_checkpoint_policy, list_operators, selective_checkpoint_wrapper) is unchanged, except that in-place ops inside a checkpointed region are no longer supported (see below).

Changed

• The implementation of xformers.ops.fmha now lives in the mslk package, which is now a dependency of xFormers. xformers.ops.fmha re-exports it, so its API is unchanged. In principle, the version of mslk must be ABI-compatible with your version of PyTorch (see its compatibility table). On CUDA, xFormers tolerates an incompatible mslk, as its default attention backends only need mslk's Python code: it switches mslk to its python-only mode, with a warning. On ROCm, some attention backends (e.g. CK) are unavailable if mslk is incompatible.
• In-place ops inside a region checkpointed with an xformers.checkpoint policy are no longer supported, as PyTorch 2.13+ rejects them. get_optimal_checkpoint_policy now warns when the policy it builds stores an in-place op; use non-in-place activations instead.
• On Windows, FlashAttention-3 is no longer available out of the box: flash-attn-3 only has Windows builds for CUDA 12.x, and mslk only for CUDA 13.x. With CUDA 12.x, installing a python-only mslk from source works around this (see the README's troubleshooting section).

Fixed

• import xformers.ops no longer crashes on PyTorch builds without distributed support, such as the NVIDIA PyTorch containers for Jetson iGPU devices, where torch.distributed.is_available() is False. The model-parallel and sequence-parallel ops, which cannot work at all there, are simply not imported, and everything else stays usable.
• Reusing the policy returned by get_optimal_checkpoint_policy across several checkpointed regions no longer raises an IndexError.
• The profiler no longer over-counts the FLOPs of the softmax@V product for non-square causal attention.
• The fused reduce-scatter sequence-parallel ops now validate their out argument correctly.

Deprecated

• The xformers.checkpoint API (checkpoint, selective_checkpoint_context_fn, get_optimal_checkpoint_policy, selective_checkpoint_wrapper, list_operators) is deprecated and will be removed in a future release; calling it now raises a FutureWarning. PyTorch provides equivalents: torch.utils.checkpoint.create_selective_checkpoint_contexts for policy-based selective checkpointing, and the torch.compile activation memory budget (torch._functorch.config.activation_memory_budget) for memory-budget-driven recomputation.

Removed

• Removed 2:4 sparsity support: xformers.ops.sparsify24, xformers.ops.sparsify24_like, xformers.ops.Sparse24Tensor, everything else in xformers.ops.sp24, and the underlying xformers::sparse24_* operators. PyTorch supports semi-structured (2:4) sparsity natively through torch.sparse.to_sparse_semi_structured and torch.sparse.SparseSemiStructuredTensor.