flash-attn-4

library4.0.0b12pypypiunverified

Flash Attention 4 is the next-generation implementation of the Flash Attention algorithm using NVIDIA CUTE (CUDA Template Engine). It provides highly optimized fused attention kernels for modern GPUs, supporting head dimensions up to 256 and various data types including FP8. Version 4.0.0b12 is in beta, with frequent releases.

ai-ml llm-agents

pip install flash-attn-4

github ↗package ↗

INSTALL

IMPORT

SIG · FLASH-ATTN-4

flash-attn-4

ai-mlpythonv4.0.0b12

harness data pending

Install & Compatibility

Where this runs

No compatibility data collected yet for this library.

Code

Verified usage

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

flash_attn_func

✓ from flash_attn_4 import flash_attn_func

✗ from flash_attn import flash_attn_func

Flash Attention 4 uses a separate PyPI package 'flash-attn-4' with module name 'flash_attn_4'. Importing from old 'flash_attn' imports Flash Attention 2/3.

flash_attn_varlen_func

✓ from flash_attn_4 import flash_attn_varlen_func

Correct import for variable-length sequences.

Basic forward pass with causal masking.

import torch from flash_attn_4 import flash_attn_func q = torch.randn(1, 4, 128, 64, device='cuda', dtype=torch.float16) k = torch.randn(1, 4, 128, 64, device='cuda', dtype=torch.float16) v = torch.randn(1, 4, 128, 64, device='cuda', dtype=torch.float16) out, lse = flash_attn_func(q, k, v, causal=True) print(out.shape)

Debug

Known issues

breakingFlash Attention 4 is a completely new implementation using CUTE. The API has changed; functions like flash_attn_func now return a tuple (out, lse) instead of just out.

fix

Update code to unpack the tuple: out, lse = flash_attn_func(...)

affects: >=4.0.0b1

gotchaThe PyPI package name is 'flash-attn-4', and the Python module is 'flash_attn_4'. Do not confuse with the old 'flash-attn' package (Flash Attention 2/3).

fix

Use 'pip install flash-attn-4' and 'import flash_attn_4'.

affects: all

gotchaFlash Attention 4 only supports CUDA GPUs with compute capability 8.0+ (Ampere, Hopper, Blackwell). It will fail on older GPUs.

fix

Check GPU compute capability via torch.cuda.get_device_capability(). Minimum 8.0 required.

affects: all

Upgrade

Version history

4.0.0b12latest on PyPI · released May 6, 2026

Audit

Dependencies

torchrequiredRequired for tensor operations and GPU support.

Agent activity

20 hits · last 30 days

node

OpenAI (training)

Amazon

Resources

githubhttps://github.com/Dao-AILab/flash-attention ↗

packagehttps://pypi.org/project/flash-attn-4/ ↗