Registry / ai-ml / flash-attn

flash-attn

JSON →
library2.8.3.post1pypypiunverified

Flash Attention is a fast and memory-efficient exact attention mechanism for deep learning models, particularly Transformers. It reorders the attention computation to reduce the number of memory accesses, making it significantly faster and less memory-intensive than standard attention. The library is currently stable at version 2.8.3, with an active beta development for version 4.0.0 which introduces new features and architectural changes. Its release cadence is driven by research advancements and performance optimizations.

pip install flash-attn --no-build-isolation
INSTALL
IMPORT
SIG · FLASH-ATTN
F
flash-attn
ai-mlpythonv2.8.3.post1
Install
—
Import
—
Disk
—
Pass rate
0/ 10
Env Coverage0 / 10
glibc
3.9–3.13
musl
3.9–3.13
Install & Compatibility
Where this runs
tested against v? · pip install
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
musl
py 3.10–3.910 runs
build_error
glibc
py 3.10–3.910 runs
build_error
Code
Verified usage

Verified import paths — ran on the pinned version, not inferred.

flash_attn_func
✓ from flash_attn import flash_attn_func
flash_attn_qkvpacked_func
✓ from flash_attn import flash_attn_qkvpacked_func
flash_attn_varlen_func
✓ from flash_attn import flash_attn_varlen_func
FlashAttention2
✓ from flash_attn.modules.mha import FlashAttention2

This quickstart demonstrates how to use `flash_attn_func` with separate query, key, and value tensors. It highlights the importance of data types (float16/bfloat16) and device placement (CUDA) for optimal performance. The `causal=True` argument is common for generative models. Ensure your `head_dim` is a multiple of 8 and ideally no more than 256.

import torch from flash_attn import flash_attn_func # Example for q, k, v as separate tensors batch_size = 2 seq_len = 128 num_heads = 8 head_dim = 64 # Must be multiple of 8, typically <= 256 dtype = torch.float16 # FlashAttention works best with float16 or bfloat16 device = 'cuda' if torch.cuda.is_available() else 'cpu' if device == 'cpu': print("Warning: FlashAttention is primarily designed for CUDA GPUs.") q = torch.randn(batch_size, seq_len, num_heads, head_dim, dtype=dtype, device=device) k = torch.randn(batch_size, seq_len, num_heads, head_dim, dtype=dtype, device=device) v = torch.randn(batch_size, seq_len, num_heads, head_dim, dtype=dtype, device=device) # Causal attention (for language models) output = flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=True) print("Output shape:", output.shape) print("Output device:", output.device)
Debug
Known issues
breakingThe API for `flash_attn_func` has changed significantly between v1, v2, and the v4 beta, including argument order, default values, and added parameters (e.g., `softmax_scale`, `dropout_p`, `causal`, different return values). Code written for v1 or early v2 will likely break on later v2 or v4.
fix
Refer to the official documentation or GitHub README for the exact function signature matching your installed version. Pay close attention to `qkv` vs `q, k, v` inputs and boolean flags.
affects: <2.0, 2.x, 4.x (beta)
gotchaFlash Attention requires a specific CUDA architecture (SM70+ for v1/v2, SM80+ for v2.2+, SM90+ for v4 beta) and specific `head_dim` values. Typically, `head_dim` must be a multiple of 8 (e.g., 64, 128, 256) and for optimal performance, should not exceed 256. Using unsupported `head_dim` or CUDA architecture will result in runtime errors or fallbacks to slower implementations.
fix
Ensure your GPU supports the required CUDA architecture. Adjust `head_dim` to be a multiple of 8. For best performance, keep `head_dim <= 256`. Check the FlashAttention GitHub for exact hardware requirements.
affects: All
gotchaInstallation can be sensitive to your PyTorch and CUDA setup. Using `pip install flash-attn` without `--no-build-isolation` can lead to `flash-attn` compiling against a different CUDA toolkit than your PyTorch installation, causing runtime errors or crashes.
fix
Always use `pip install flash-attn --no-build-isolation` when installing with CUDA extensions. Ensure your PyTorch version and CUDA toolkit are compatible as specified by PyTorch. Consider installing `ninja` first (`pip install ninja`) for smoother compilation.
affects: All
breakingThe FlashAttention v4 beta introduces new APIs and internal architecture changes. Code written for v2.x is NOT directly compatible with the v4 beta, and vice versa. Key changes include a redesigned `FlashAttention2` module and `flash_attn_func` with updated arguments to support new features like dynamic sequence lengths.
fix
If migrating to or from v4 beta, review the specific release notes and documentation for v4. Adapt your code to the new API, especially for module initialization and function calls.
affects: 4.0.0.betaX
Upgrade
Version history
2.8.3.post1latest on PyPI · released Jun 11, 2026
Audit
Dependencies
torchrequiredCore deep learning framework
cudaoptionalRequired for GPU acceleration and primary performance benefits
ninjaoptionalOften helps with C++/CUDA extension compilation if `pip` struggles
Agent activity
36 hits · last 30 days
node
32
OpenAI (training)
1
Resources